Update README.md

README.md

@@ -154,23 +154,40 @@ This model uses a specific chat format for optimal performance.
 ## Usage with HuggingFace transformers 
 The model can be used with HuggingFace's `transformers` library:
 ```python
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model = AutoModelForCausalLM.from_pretrained("OuteAI/Lite-Mistral-150M-v2-Instruct")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = AutoModelForCausalLM.from_pretrained("OuteAI/Lite-Mistral-150M-v2-Instruct").to(device)
 tokenizer = AutoTokenizer.from_pretrained("OuteAI/Lite-Mistral-150M-v2-Instruct")
 
-def generate_response(message):
-    # Encode the formatted message as input ids
-    input_ids = tokenizer.encode(f"<s>user\n{message}</s>", return_tensors="pt")
-    output = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
+def generate_response(message: str, temperature: float = 0.4, repetition_penalty: float = 1.1) -> str:
+    # Apply the chat template and convert to PyTorch tensors
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": message}
+    ]
+    input_ids = tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt"
+    ).to(device)
+
+    # Generate the response
+    output = model.generate(
+        input_ids,
+        max_length=512,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        do_sample=True
+    ) 
 
     # Decode the generated output
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    
     return generated_text
 
-message = "What is the capital of Spain?"
+message = "I'd like to learn about language models. Can you break down the concept for me?"
 response = generate_response(message)
+print(response)
 ```
 
 ## Risk Disclaimer