Update README.md

README.md

@@ -44,23 +44,33 @@ model = AutoModelForCausalLM.from_pretrained("AstroMLab/astrollama-2-7b-chat_aic
 
 # Function to generate a response
 def generate_response(prompt, max_length=512):
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_length)
+    full_prompt = f"###Human: {prompt}\n\n###Assistant:"
+    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=max_length)
     inputs = inputs.to(model.device)
     
     # Generate a response
     with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=max_length, num_return_sequences=1, do_sample=True)
+        outputs = model.generate(
+            **inputs, 
+            max_length=max_length, 
+            num_return_sequences=1, 
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.encode("###Human:", add_special_tokens=False)[0]
+        )
     
     # Decode and return the response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    
+    # Extract only the Assistant's response
+    assistant_response = response.split("###Assistant:")[-1].strip()
+    return assistant_response
 
-# Example conversation
+# Example usage
 user_input = "What are the main components of a galaxy?"
-prompt = f"Human: {user_input}\n\nAssistant:"
-
-response = generate_response(prompt)
-print(response)
+response = generate_response(user_input)
+print(f"Human: {user_input}")
+print(f"Assistant: {response}")
 ```
 
 ## Model Limitations and Biases