Update README.md

README.md

@@ -31,38 +31,28 @@ Due to its size, the model runs very well on CPUs.
 
 ```python
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-from peft import PeftModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-base_model_id = "unsloth/Llama-3.2-1B-Instruct"
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
+model_id = "LeonardPuettmann/LlaMaestra-3.2-1B-Instruct-v0.1-4bit"
 
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id,  # Mistral, same as before
-    quantization_config=bnb_config,  # Same quantization config as before
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
     device_map="auto",
     trust_remote_code=True,
 )
 
-tokenizer = AutoTokenizer.from_pretrained(base_model_id, add_bos_token=True, trust_remote_code=True)
-
-ft_model = PeftModel.from_pretrained(base_model, "LeonardPuettmann/LlaMaestra-3.2-1B-Instruct-v0.1-4bit")
+tokenizer = AutoTokenizer.from_pretrained(model_id, add_bos_token=True, trust_remote_code=True)
 
 row_json = [
     {"role": "system", "content": "Your job is to return translations for sentences or words from either Italian to English or English to Italian."},
-    {"role": "user", "content": "Scontri a Bologna, la destra lancia l'offensiva contro i centri sociali."}
+    {"role": "user", "content": "Do you sell tickets for the bus?"},
 ]
 
 prompt =  tokenizer.apply_chat_template(row_json, tokenize=False)
 model_input = tokenizer(prompt, return_tensors="pt").to("cuda")
 
 with torch.no_grad():
-    print(tokenizer.decode(ft_model.generate(**model_input, max_new_tokens=1024)[0]))
+    print(tokenizer.decode(model.generate(**model_input, max_new_tokens=1024)[0]))
 ```
 
 ## Data used