Update README.md

README.md

@@ -19,11 +19,11 @@ ChatGLM3-6B中文纠错LoRA模型
 
 The overall performance of shibing624/chatglm3-6b-csc-chinese-lora on CSC **test**:
 
-|prefix|input_text|target_text|pred|
-|:-- |:--- |:--- |:-- |
-|对下面文本纠错：|少先队员因该为老人让坐。|少先队员应该为老人让座。|少先队员应该为老人让座。|
+|prefix|input_text|pred|
+|:-- |:--- |:--- |
+|对下面文本纠错：|少先队员因该为老人让坐。|少先队员应该为老人让座。|
 
-在CSC测试集上生成结果纠错准确率高，由于是基于ChatGLM3-6B模型，结果常常能带给人惊喜，不仅能纠错，还带有句子润色和改写功能。
+在CSC测试集上生成结果纠错准确率高，由于是基于[THUDM/chatglm3-6b](https://huggingface.co/THUDM/chatglm3-6b)模型，结果常常能带给人惊喜，不仅能纠错，还带有句子润色和改写功能。
 
 
 ## Usage
@@ -53,21 +53,35 @@ pip install transformers
 ```
 
 ```python
-import sys
-from peft import PeftModel
-from transformers import AutoModel, AutoTokenizer
+import os
 
-sys.path.append('..')
+import torch
+from peft import PeftModel
+from transformers import AutoTokenizer, AutoModel
 
-model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True, device_map='auto')
-model = PeftModel.from_pretrained(model, "shibing624/chatglm3-6b-csc-chinese-lora")
-model = model.half().cuda()  # fp16
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True)
+model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True).half().cuda()
+model = PeftModel.from_pretrained(model, "shibing624/chatglm3-6b-csc-chinese-lora")
+
+sents = ['对下面文本纠错\n\n少先队员因该为老人让坐。',
+         '对下面文本纠错\n\n下个星期，我跟我朋唷打算去法国玩儿。']
+
+
+def get_prompt(user_query):
+    vicuna_prompt = "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: {query} ASSISTANT:"
+    return vicuna_prompt.format(query=user_query)
+
 
-sents = ['对下面中文拼写纠错：\n少先队员因该为老人让坐。',
-         '对下面中文拼写纠错：\n下个星期，我跟我朋唷打算去法国玩儿。']
 for s in sents:
-    response = model.chat(tokenizer, s, max_length=128, eos_token_id=tokenizer.eos_token_id)
+    q = get_prompt(s)
+    input_ids = tokenizer(q).input_ids
+    generation_kwargs = dict(max_new_tokens=128, do_sample=True, temperature=0.8)
+    outputs = model.generate(input_ids=torch.as_tensor([input_ids]).to('cuda'), **generation_kwargs)
+    output_tensor = outputs[0][len(input_ids):]
+    response = tokenizer.decode(output_tensor, skip_special_tokens=True)
     print(response)
 ```