Update README.md

README.md

@@ -33,21 +33,31 @@ This model can answer information about different excplicit ideas in medicine
 
 ```python
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-torch.set_default_device("cuda")
+base_model_id = "microsoft/phi-2"
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_id,  # Phi2, same as before
+    device_map="auto",
+    trust_remote_code=True,
+    load_in_8bit=True,
+    torch_dtype=torch.float16,
+)
 
-model = AutoModelForCausalLM.from_pretrained("segestic/phi2_medical_seg", torch_dtype="auto", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("segestic/phi2_medical_seg", trust_remote_code=True)
+eval_tokenizer = AutoTokenizer.from_pretrained(base_model_id, add_bos_token=True, trust_remote_code=True, use_fast=False)
+eval_tokenizer.pad_token = tokenizer.eos_token
 
-inputs = tokenizer('''def print_prime(n):
-   """
-   What is Medcine?
-   """''', return_tensors="pt", return_attention_mask=False)
+from peft import PeftModel
 
-outputs = model.generate(**inputs, max_length=200)
-text = tokenizer.batch_decode(outputs)[0]
-print(text)
+adapter_model_id = "segestic/phi2_medical_seg"
+ft_model = PeftModel.from_pretrained(base_model, adapter_model_id)
+
+eval_prompt = "What is medicine?"
+model_input = eval_tokenizer(eval_prompt, return_tensors="pt").to("cuda")
+
+ft_model.eval()
+with torch.no_grad():
+    print(eval_tokenizer.decode(ft_model.generate(**model_input, max_new_tokens=100, repetition_penalty=1.11)[0], skip_special_tokens=True))
 ```