from config import *
from utils import *
from model import make_model
import torch
from torch.nn.utils.rnn import pad_sequence

if __name__ == '__main__':
    en_id2vocab, en_vocab2id  = get_vocab('en')
    zh_id2vocab, zh_vocab2id  = get_vocab('zh')

    SRC_VOCAB_SIZE = len(en_id2vocab) 
    TGT_VOCAB_SIZE = len(zh_id2vocab) 

    model = make_model(SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, D_MODEL, N_HEAD, D_FF, N, DROPOUT).to(device)
    model.load_state_dict(
        torch.load("./transformer_code/params/best.pt",
                   map_location=device))
    
    model.eval()
    # "其一是不久前做出的对本可用于住宅开发的未利用土地征税的决策。"
    # "投资不足最终会削弱经济增长并减少就业机会。"
    # "其真正的目标是引导构建一个包容性和合作的一体化过程，令所有参与者都能从中受益。"

    # test_clean.json中的数据
    texts = [
        "I like playing basketball",
        "He is a doctor",
        "I am a student",
        "What is particularly problematic in the EU is the clash between radical politics and mainstream governance.",
        # "欧洲问题的严重之处在于激进政治与主流治理之间的碰撞。"
        "Such a resounding collapse of the Nuclear Non-Proliferation Treaty would pose a formidable challenge to the world community, particularly to the Middle East.",
        # "《核不扩散条约》如此引人瞩目的崩溃会给国际社会，特别是中东国家带来可怕的挑战。"
        "It certainly would deal a devastating blow to the EU’s credibility, both in its member states and abroad.",
        # "这无疑会在欧盟成员国及海外对欧盟的信誉造成毁灭性的打击。"
        "Its real goal is to lead an inclusive and cooperative integration process that benefits all participants.",
        # "其真正的目标是引导构建一个包容性和合作的一体化过程，令所有参与者都能从中受益。"
    ]
    
    #把词转成对应的索引
    batch_src_token = [[en_vocab2id.get(v.lower(), UNK_ID) for v in divided_en(text)] for text in texts]
    batch_src = [torch.LongTensor(src) for src in batch_src_token]
    #每句话加上padding
    src_x = pad_sequence(batch_src, True, PAD_ID)
    src_mask = get_padding_mask(src_x, PAD_ID)
    
    prob_sent = batch_greedy_decode(model, src_x, src_mask)

    print(prob_sent)