Update README.md

README.md

@@ -26,19 +26,21 @@ license_link: LICENSE
 可通过以下代码加载 Index-1.9B-Chat 模型来进行对话：
 
 ```python
-import torch
 import argparse
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, pipeline
 
 # 注意！目录不能含有"."，可以替换成"_"
 parser = argparse.ArgumentParser()
-parser.add_argument('--model_path', default="", type=str, help="")
+parser.add_argument('--model_path', default="IndexTeam/Index-1.9B-Chat", type=str, help="")
+parser.add_argument('--device', default="cpu", type=str, help="") # also could be "cuda" or "mps" for Apple silicon
 args = parser.parse_args()
 
 tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(args.model_path, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map='auto')
-model = model.eval()
-print('model loaded', args.model_path, model.device)
+generator = pipeline("text-generation",
+                    model=args.model_path,
+                    tokenizer=tokenizer, trust_remote_code=True, 
+                    device=args.device)
+
 
 system_message = "你是由哔哩哔哩自主研发的大语言模型，名为“Index”。你能够根据用户传入的信息，帮助用户完成指定的任务，并生成恰当的、符合要求的回复。"
 query = "续写 天不生我金坷垃"
@@ -46,15 +48,8 @@ model_input = []
 model_input.append({"role": "system", "content": system_message})
 model_input.append({"role": "user", "content": query})
 
-inputs = tokenizer.apply_chat_template(model_input, tokenize=False, add_generation_prompt=False)
-input_ids = tokenizer.encode(inputs, return_tensors="pt").to(model.device)
-history_outputs = model.generate(input_ids, max_new_tokens=300, top_k=5, top_p=0.8, temperature=0.3, repetition_penalty=1.1, do_sample=True)
-
-# 删除</s>
-if history_outputs[0][-1] == 2:
-    history_outputs = history_outputs[:, :-1]
+model_output = generator(model_input, max_new_tokens=300, top_k=5, top_p=0.8, temperature=0.3, repetition_penalty=1.1, do_sample=True)
 
-outputs = history_outputs[0][len(input_ids[0]):]
 print('User:', query)
-print('\nModel:', tokenizer.decode(outputs))
+print('Model:', model_output)
 ```