翻译的两种实现方式——基于transformers_opus-mt-en-zh

1.基于pipeline 

(可加参数 device,默认为-1,使用cpu，如果非负，则表示指定哪块gpu)

from transformers import (
     AutoTokenizer,
     AutoModelForSeq2SeqLM,
     pipeline
)
 
text = "从时间上看，中国空间站的建造比国际空间站晚20多年。"
 
tokenizer = AutoTokenizer.from_pretrained("./Helsinki-NLP/opus-mt-zh-en")
model = AutoModelForSeq2SeqLM.from_pretrained("./Helsinki-NLP/opus-mt-zh-en")
 
tokenizer_back_translate = AutoTokenizer.from_pretrained("./Helsinki-NLP/opus-mt-en-zh")
model_back_translate = AutoModelForSeq2SeqLM.from_pretrained("./Helsinki-NLP/opus-mt-en-zh")
 
zh2en = pipeline("translation_zh_to_en", model=model, tokenizer=tokenizer)
en2zh = pipeline("translation_en_to_zh", model=model_back_translate, tokenizer=tokenizer_back_translate)
print("tran", zh2en(text[:5])[0]['translation_text'])
print("tran_back", en2zh(zh2en(text[:5])[0]['translation_text'], max_length=510)[0]['translation_text'])

2.逐步实现

batch = tokenizer.prepare_seq2seq_batch(src_texts=[text], return_tensors='pt', max_length=512)
# Perform the translation and decode the output
translation = model.generate(**batch)
result = tokenizer.batch_decode(translation, skip_special_tokens=True)
print("tran", result)
batch_back_translate = tokenizer_back_translate.prepare_seq2seq_batch(src_texts=result, return_tensors='pt', max_length=512)
# Perform the translation and decode the output
translation_back_translate = model_back_translate.generate(**batch_back_translate)
result = tokenizer_back_translate.batch_decode(translation_back_translate, skip_special_tokens=True)
print("tran_back", result)