Update README.md

README.md

@@ -30,24 +30,22 @@ The code used to generate cosmosage_v2 is available at https://github.com/tijmen
 After downloading cosmosage_v2, the following example code can be used to ask questions:
 
 ```python
-path_to_model = 'cosmosage_v2/'
+model_path = "models/cosmosage_v2/"
 
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-device = "cuda"
-model = AutoModelForCausalLM.from_pretrained(path_to_model).to(device)
-tokenizer = AutoTokenizer.from_pretrained(path_to_model)
-def ask_cosmosage(question):
-    input_ids = torch.cat([
-    tokenizer.encode("You are cosmosage, an AI programmed to be a cosmology expert. You answer the USER's question clearly in long form, always providing context. When appropriate, provide a reference.", return_tensors="pt"),
-    torch.tensor([[28705]]),
-    tokenizer.encode("USER:", add_special_tokens=False, return_tensors="pt"),
-    tokenizer.encode(question, add_special_tokens=False, return_tensors="pt"),
-    torch.tensor([[28705]]),
-    tokenizer.encode("ASSISTANT:", add_special_tokens=False, return_tensors="pt")
-    ], dim=-1).to(device)
-    generated_ids = model.generate(input_ids, max_length=input_ids.shape[1] + 1000, do_sample=True, temperature=0.4)
-    return tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+device = torch.device("cuda")
+model = AutoModelForCausalLM.from_pretrained(model_path).to(device, dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+def ask_cosmosage(question, answer_start=''):
+    prompt = f"You are cosmosage, an AI programmed to be a cosmology expert. You answer the USER's question clearly in long form, always providing context. When appropriate, provide a reference.USER: {question}ASSISTANT: {answer_start}"
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    generated_ids = model.generate(input_ids, max_length=1024, do_sample=True, temperature=0.7, top_k=None, pad_token_id=tokenizer.eos_token_id)
+    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    answer = generated_text.split("ASSISTANT:")[-1]
+    return answer
 ```
 
 ## Comparison to cosmosage_v1