Update README.md

README.md

@@ -34,6 +34,7 @@ Versions:
   - Hardware: 1x RTX 3090
 
 ```python
+import torch
 from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 
@@ -43,10 +44,19 @@ config = PeftConfig.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     config.base_model_name_or_path,
     load_in_8bit=True,
+    torch_dtype=torch.float16,
     device_map="auto"
 )
-model = PeftModel.from_pretrained(model, MODEL_NAME, device_map={'': 0})
+model = PeftModel.from_pretrained(model, MODEL_NAME, torch_dtype=torch.float16)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model.eval()
+
+generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
+data = tokenizer(["Вопрос: Почему трава зеленая?\n\nВыход:"], return_tensors="pt")
+data = {k: v.to(model.device) for k, v in data.items() if k in ("input_ids", "attention_mask")}
+output_ids = model.generate(**data, generation_config=generation_config)[0]
+print(tokenizer.decode(output_ids, skip_special_tokens=True))
+
 ```
 
 Examples: