Update README.md

README.md

@@ -104,21 +104,20 @@ First, you pass your input through the transformer model, then you get the gener
 
 Install package:
 ```
-pip install transformers 
+pip install transformers -U
 ```
 
 ```python
 import sys
+import torch
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 
 
 model = AutoModelForCausalLM.from_pretrained("baichuan-inc/Baichuan-13B-Chat", device_map='auto', trust_remote_code=True)
 model.generation_config = GenerationConfig.from_pretrained("baichuan-inc/Baichuan-13B-Chat", trust_remote_code=True)
-
 tokenizer = AutoTokenizer.from_pretrained("baichuan-inc/Baichuan-13B-Chat", trust_remote_code=True)
 model = PeftModel.from_pretrained(model, "shibing624/vicuna-baichuan-13b-chat-lora")
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def generate_prompt(instruction):
@@ -131,7 +130,7 @@ for s in sents:
     inputs = tokenizer(q, return_tensors="pt")
     inputs = inputs.to(device=device)
 
-    generate_ids = ref_model.generate(
+    generate_ids = model.generate(
         **inputs,
         max_new_tokens=120, 
     )