from transformers import BartForConditionalGeneration, BertTokenizer
import torch

model_path = "yanghh7/Alirector-bart"

tokenizer = BertTokenizer.from_pretrained(model_path)
    
model = BartForConditionalGeneration.from_pretrained(model_path).cuda()

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

    model_inputs = tokenizer(
        source,
        return_tensors='pt',
        return_token_type_ids=False,
    ).to(model.device)
    
    with torch.no_grad():
        output = model.generate(**model_inputs,)
    response = tokenizer.batch_decode(output.detach().cpu(), skip_special_tokens=True)[0]
    response = response.replace(' ', '')
    print(response)

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",
}