import ctranslate2
from split import split_string

translator_zh2en = ctranslate2.Translator("zh-en_model/", device="cpu")##路径
translator2_zh2en = ctranslate2.Translator("zh2en_cmodel/", device="cpu")##路径
translator_en2zh = ctranslate2.Translator("en-zh_model/", device="cpu")##路径
translator2_en2zh = ctranslate2.Translator("en2zh_cmodel", device="cpu")##路径

def translate(input_tokens, input_tokens2, mode):
    
    source = split_string(input_tokens)
    lenth = len(source)

    source2 = split_string(input_tokens2)
    lenth2 = len(source2)

    if mode == "汉译英" :
        results = translator_zh2en.translate_batch(source)##翻译的分词分句
        results2 = translator2_zh2en.translate_batch(source2)##翻译的分词分句
    else :
        results = translator_en2zh.translate_batch(source)##翻译的分词分句
        results2 = translator2_en2zh.translate_batch(source2)##翻译的分词分句

    target = []
    target2 = []
    
    for i in range(0, lenth, 1):
        target = target + results[i].hypotheses[0]
    for i in range(0, lenth2, 1):
        target2 = target2 + results2[i].hypotheses[0]

    #print(results[0].hypotheses[0])##results[0]为第0句，hypotheses[0]保持0
    ##print(results[1].hypotheses[0])
    #return results[0].hypotheses[0]
    return target,target2