Views
No views yet
1import transformers
2import torch
3
4model_path = "SNOWTEAM/medico-mistral"
5model = AutoModelForCausalLM.from_pretrained(
6 model_path,device_map="auto",
7 max_memory=max_memory_mapping,
8 torch_dtype=torch.float16,
9)
10tokenizer = AutoTokenizer.from_pretrained("SNOWTEAM/medico-mistral")
11input_text = ""
12input_ids = tokenizer(input_text, return_tensors="pt").input_ids
13output_ids = model.generate(input_ids=input_ids.cuda(),
14 max_new_tokens=300,
15 pad_token_id=tokenizer.eos_token_id,)
16output_text = tokenizer.batch_decode(output_ids[:, input_ids.shape[1]:],skip_special_tokens=True)[0]
17print(output_text)