Update README.md

README.md

@@ -17,8 +17,7 @@ The Mistral-7B-v0.1 model is a transformer-based model that can handle a variety
 How to utilize my Mistral for Italian text generation
 
 ```python
-import transformers
-from transformers import TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -29,14 +28,18 @@ model.to(device)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
 def generate_answer(prompt):
-    encoded_input = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], return_tensors="pt").to(device)
-    generated_ids = model.generate(**encoded_input, max_new_tokens=200, do_sample=True, temperature=0.001, eos_token_id=tokenizer.eos_token_id)
-    answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
-    return answer
-
-prompt = "Se un mattone pesa 1kg più metà di se stesso, quanto pesa il mattone? Rispondi impostando l'equazione matematica"
-print(generate_answer(prompt))
-
+    messages = [
+        {"role": "user", "content": prompt},
+    ]
+    model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(device)
+    generated_ids = model.generate(model_inputs, max_new_tokens=200, do_sample=True,
+                                          temperature=0.001, eos_token_id=tokenizer.eos_token_id)
+    decoded = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+    return decoded[0]
+
+prompt = "Come si apre un file json in python?"
+answer = generate_answer(prompt)
+print(answer)
 ```
 ---
 ## Developer