Update README.md

README.md

@@ -35,29 +35,20 @@ import transformers
 from transformers import TextStreamer
 import torch
 
-model_name = "DeepMount00/Mistral-Ita-7b"
+MODEL_NAME = "DeepMount00/Mistral-Ita-7b"
 
-tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
-model = transformers.LlamaForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto").eval()
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16).eval()
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-def stream(user_prompt):
-    runtimeFlag = "cuda:0"
-    system_prompt = ''
-    B_INST, E_INST = "<s> [INST]", "[/INST]"
-    prompt = f"{system_prompt}{B_INST}{user_prompt.strip()}\n{E_INST}"
-    inputs = tokenizer([prompt], return_tensors="pt").to(runtimeFlag)
-    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    _ = model.generate(**inputs, streamer=streamer, max_new_tokens=300, temperature=0.0001,
-                        repetition_penalty=1.2, eos_token_id=2, do_sample=True, num_return_sequences=1)
+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
 
-domanda = """Scrivi una funzione python che moltiplica per 2 tutti i valori della lista:"""
-contesto = """
-[-5, 10, 15, 20, 25, 30, 35]
-"""
+prompt = "Se un mattone pesa 1kg più metà di se stesso, quanto pesa il mattone? Rispondi impostando l'equazione matematica"
+print(generate_answer(prompt))
 
-prompt = domanda + "\n" + contesto
-
-stream(prompt)
 ```
 ---
 ## Developer