from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

base_model_name = "baichuan-inc/Baichuan2-7B-Base"
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype="auto",
    device_map="cuda:0"
)
lora_model_name = "yanghh7/Alirector-baichuan2-7b-lora"
model = PeftModel.from_pretrained(
    base_model,
    lora_model_name
)
tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
template = "对输入句子进行语法纠错，并输出正确的句子。\nInput:{source}\nOutput:"

while True:
    source = input("输入句子：")

    prompt = template.format(source=source)
    model_inputs = tokenizer(
        [prompt],
        return_tensors="pt",
    ).to(model.device)

    with torch.no_grad():
        output = model.generate(**model_inputs)
        
    response = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
    print(response.split('\nOutput:')[-1])

Citation

@inproceedings{yang-quan-2024-alirector,
    title = "Alirector: Alignment-Enhanced {C}hinese Grammatical Error Corrector",
    author = "Yang, Haihui and Quan, Xiaojun",
    booktitle = "Findings of the Association for Computational Linguistics: ACL 2024",
    year = "2024",
}