Update README.md

README.md

@@ -114,11 +114,11 @@ 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 = AutoModelForCausalLM.from_pretrained("baichuan-inc/Baichuan-13B-Chat", device_map='auto', torch_dtype=torch.float16, 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"
+device = torch.device(0) if torch.cuda.is_available() else torch.device("cpu")
 
 def generate_prompt(instruction):
     return f"""A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.USER: {instruction} ASSISTANT: """
@@ -128,11 +128,11 @@ sents = ['一岁宝宝发烧能吃啥药', "who are you?"]
 for s in sents:
     q = generate_prompt(s)
     inputs = tokenizer(q, return_tensors="pt")
-    inputs = inputs.to(device=device)
+    inputs = inputs.to(device)
 
     generate_ids = model.generate(
         **inputs,
-        max_new_tokens=120, 
+        max_new_tokens=512, 
     )
 
     output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0]
@@ -154,6 +154,7 @@ vicuna-baichuan-13b-chat-lora
     └── adapter_model.bin
 ```
 
+- Inference GPU: 27G
 
 ### Inference Examples