Knight is a text style transfer model for knight-errant style. This model is for Chinese Knight-errant style transfer.
1
2#inference
3from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
4
5model = MBartForConditionalGeneration.from_pretrained("Anonymous-TST/knight-errant-TST-zh")
6tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50", src_lang="zh_CN", tgt_lang="zh_CN")
7
8model.cuda()
9model.eval()
10
11article_1 = "jinyong: 接下来会发生什么？"
12batch = tokenizer(article_1, return_tensors="pt",return_token_type_ids=False, truncation=True, max_length=64, padding=True).to('cuda')
13translated_tokens = model.generate(**batch,max_length=64)
14decoded = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=True)
15print(decoded)
16# 欲知后事如何?
