Update README.md

README.md

@@ -12,25 +12,40 @@ pip install transformers
 ```
 
 ```python
-# pip install transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
-checkpoint = "twnlp/ChineseErrorCorrector2-7B"
 
-device = "cuda" # for GPU usage or "cpu" for CPU usage
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
+model_name = "twnlp/ChineseErrorCorrector2-7B"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+prompt = "你是一个文本纠错专家，纠正输入句子中的语法错误，并输出正确的句子，输入句子为："
+text_input = "少先队员因该为老人让坐。"
+messages = [
+    {"role": "user", "content": prompt + text_input}
+]
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+generated_ids = model.generate(
+    **model_inputs,
+    max_new_tokens=512
+)
+generated_ids = [
+    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+]
+
+response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+print(response)
 
-input_content = "你是一个文本纠错专家，纠正输入句子中的语法错误，并输出正确的句子，输入句子为：\n少先队员因该为老人让坐。"
-
-messages = [{"role": "user", "content": input_content}]
-input_text=tokenizer.apply_chat_template(messages, tokenize=False)
-
-print(input_text)
-
-inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
-outputs = model.generate(inputs, max_new_tokens=1024, temperature=0, do_sample=False, repetition_penalty=1.08)
-
-print(tokenizer.decode(outputs[0]))
 ```
 
 output: