Update README.md

README.md

@@ -61,18 +61,19 @@ Performance on multi-sentence texts is not guaranteed, please be aware.
 
 ```python
 # pip install bitsandbytes transformers peft torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftConfig, PeftModel
 import torch
 
 config = PeftConfig.from_pretrained("lang-uk/dragoman")
 quant_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=float16,
+    bnb_4bit_compute_dtype=torch.float16,
     bnb_4bit_use_double_quant=False,
 )
 
-model = MistralForCausalLM.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     "mistralai/Mistral-7B-v0.1", quantization_config=quant_config
 )
 model = PeftModel.from_pretrained(model, "lang-uk/dragoman").to("cuda")
@@ -83,7 +84,7 @@ tokenizer = AutoTokenizer.from_pretrained(
 input_text = "[INST] who holds this neighborhood? [/INST]" # model input should adhere to this format
 input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
 
-outputs = model.generate(**input_ids)
+outputs = model.generate(**input_ids, num_beams=10)
 print(tokenizer.decode(outputs[0]))
 ```