Update README.md

README.md

@@ -13,7 +13,11 @@ import torch
 import transformers as tr
 
 tokenizer = tr.AutoTokenizer.from_pretrained("riccorl/Modello-Italia-9B")
-model = tr.AutoModelForCausalLM.from_pretrained("riccorl/Modello-Italia-9B", device_map="auto", torch_dtype=torch.bfloat16)
+model = tr.AutoModelForCausalLM.from_pretrained(
+  "riccorl/Modello-Italia-9B",
+  device_map="auto",
+  torch_dtype=torch.bfloat16
+)
 
 MY_SYSTEM_PROMPT_SHORT = (
   "Tu sei Modello Italia, un modello di linguaggio naturale addestrato da iGenius."
@@ -23,7 +27,9 @@ messages = [
   {"role": "system", "content": MY_SYSTEM_PROMPT_SHORT},
   {"role": "user", "content": prompt},
 ]
-tokenized_chat = hf_tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+tokenized_chat = tokenizer.apply_chat_template(
+  messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
+)
 out = model.generate(tokenized_chat.to("cuda"), max_new_tokens=200)
 ```