Update README.md

README.md

@@ -25,7 +25,7 @@ model_name = "DISCOStudy/PhenoGemini"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
 # Load base model
-base = AutoModelForCausalLM.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype="auto",
     device_map="auto",
@@ -39,8 +39,15 @@ This model does not rank over the full vocabulary. It first masks non-gene token
 
 ```python
 with torch.no_grad():
-    prompt = "[PUT THE PROMPT GENERATED FROM https://phenogemini.org]"
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    messages = [
+    {"role": "user", "content": "[PUT THE PROMPT GENERATED FROM https://phenogemini.org]"}
+]
+    text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+)
+    inputs = tokenizer([text], return_tensors="pt").to(model.device)
     outputs = model(inputs["input_ids"])
     # Take logits for the last token only
     logits = outputs.logits[0][-1]