from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer, pipeline)
import time
import os
 
class BackTranslate:
    def __init__(self, pretrained_dir):
        # 加载zh-to-en模型
        tokenizer = AutoTokenizer.from_pretrained(os.path.join(pretrained_dir, "Helsinki-NLP/opus-mt-zh-en"))
        model = AutoModelForSeq2SeqLM.from_pretrained(os.path.join(pretrained_dir, "Helsinki-NLP/opus-mt-zh-en"))
        # 加载en-to-zh模型
        tokenizer_back_translate = AutoTokenizer.from_pretrained(os.path.join(pretrained_dir, "Helsinki-NLP/opus-mt-en-zh"))
        model_back_translate = AutoModelForSeq2SeqLM.from_pretrained(os.path.join(pretrained_dir, "Helsinki-NLP/opus-mt-en-zh"))
        # 创建zh2en和en2zh管道
        self.zh2en = pipeline("translation_zh_to_en", model=model, tokenizer=tokenizer)
        self.en2zh = pipeline("translation_en_to_zh", model=model_back_translate, tokenizer=tokenizer_back_translate)
 
    def back_translate(self, text):
        """ 回译 """
        # 将输入文本翻译成英文
        text_en = self.zh2en(text, max_length=510)[0]["translation_text"]
        print("text_en:", text_en)
        # 将英文翻译回中文
        text_back = self.en2zh(text_en, max_length=510)[0]["translation_text"]
        print("text_back:", text_back)
        # 返回回译后的文本
        return text_back
 
if __name__ == '__main__':
    pretrained_dir = "D:/soft_install/dataset/bert-model/translate"
    bt = BackTranslate(pretrained_dir)
    datas = [{"text": "平乐县，古称昭州，隶属于广西壮族自治区桂林市，位于广西东北部，桂林市东南部，东临钟山县，南接昭平，西北毗邻阳朔，北连恭城，总面积1919.34平方公里。"},
     {"text": "平乐县主要旅游景点有榕津千年古榕、冷水石景苑、仙家温泉、桂江风景区、漓江风景区等，平乐县为漓江分界点，平乐以北称漓江，以南称桂江，是著名的大桂林旅游区之一。"},
     {"text": "印岭玲珑，昭水晶莹，环绕我平中。青年的乐园，多士受陶熔。生活自觉自治，学习自发自动。五育并重，手脑并用。迎接新潮流，建设新平中"},
     {"text": "桂林山水甲天下, 阳朔山水甲桂林"}
    ]
    time_start = time.time()
    for da in datas:
        text = da.get("text", "")
        bt.back_translate(text)
    time_total = time.time() - time_start
    print("time_total:{}".format(time_total))
    print("time_per:{}".format(time_total / len(datas)))