# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + gradient={"editing": false, "execution_count": 4, "id": "boolean-diving", "kernelId": "583276d5-8a0a-4f47-9ffe-672fe9cb301f"}
# %%capture
# !pip install transformers
# !pip install pypinyin
# !pip install jieba
# !pip install paddlepaddle

# + gradient={"editing": false, "execution_count": 1, "id": "metropolitan-times", "kernelId": "583276d5-8a0a-4f47-9ffe-672fe9cb301f"}
import sys
sys.path.append("../")

import re,time,json,pickle
from collections import defaultdict
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import torch.nn as nn

from transformers import (BertTokenizer,BertConfig,BertModel)
from model.fusion_transformer import Fusion_Anchi_Trans_Decoder, Fusion_Anchi_Transformer, Anchi_Decoder,Anchi_Transformer
from model.fusionDataset import FusionDataset
from utils.generate_couplet import greedy_decode,beam_search_decode

config = BertConfig.from_pretrained('../AnchiBERT')
tokenizer = BertTokenizer.from_pretrained('../AnchiBERT')
Anchibert = BertModel.from_pretrained('../AnchiBERT',config=config)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# + [markdown] gradient={"editing": false, "id": "stable-checkout", "kernelId": "583276d5-8a0a-4f47-9ffe-672fe9cb301f"}
# ### Load Necessary preproceeded Data

# + gradient={"editing": false, "execution_count": 2, "id": "funky-measurement", "kernelId": "583276d5-8a0a-4f47-9ffe-672fe9cb301f"}
with open('../pretrained_model/char_map.json','r') as f:
    ix2glyph = defaultdict(lambda : '_')
    ix2glyph[0] = '[PAD]'
    glyph2ix = defaultdict(lambda : 1)
    glyph2ix.update({'[CLS]':0,'[SEP]':0,'[PAD]':0})
    for i, k in enumerate(json.load(f).keys(),2):
        glyph2ix[k] = i
        ix2glyph[i] = k
with open('../pretrained_model/pinyin_map.json','r') as f:
    pinyin2ix = defaultdict(lambda : 1)
    pinyin2ix.update({'[CLS]':0,'[SEP]':0,'[PAD]':0})
    for i,k in enumerate(json.load(f).keys(),2):
        pinyin2ix[k] = i
with open('../pretrained_model/pos_tags.json','r') as f:
    pos2ix = defaultdict(lambda : 0)
    pos2ix.update(json.load(f))
# -

# # Decoder Section

with open("../couplet/test/in.txt",encoding='utf8') as f:
    te_in =  [row.strip().split() for row in f.readlines()]
# train 下联  
with open("../couplet/test/out.txt",encoding='utf8') as f:
    te_out = [row.strip().split() for row in f.readlines()]

# +
config = { # for Fusion_Anchi_Trans_Decoder
        'max_position_embeddings':50,
        'hidden_size':768,
        'font_weight_path':'../pretrained_model/glyph_weight.npy',
        'pinyin_embed_dim':30, # trainable
        'pinyin_path':'../pretrained_model/pinyin_map.json',
        'tag_size':30,
        'tag_emb_dim':10, # trainable
        'layer_norm_eps':1e-12,
        'hidden_dropout':0.1,
        'nhead':12,
        'num_layers':6 , #6, trainable
        'output_dim':9110,# fixed use glyph dim as output
        'device':device,
    }


# <model_name>_<optim>_<batch_num>_<lr>_<epoch>_<pinyin_embed_dim>_<tag_emb_dim>_<encoder layer>_<decoder layer>_<train_data_size>
name = 'fu_anchi_de_Adam_128_0001_60_6_30_10_51k'
model = Fusion_Anchi_Trans_Decoder(config)
model.load_state_dict(torch.load(f'../result/{name}.pt'))

# +
idx = 10

predict = greedy_decode(model=model,
                      bert=Anchibert,
                      tokenizer=tokenizer,
                      sent=te_in[idx],
                      glyph2ix=glyph2ix,
                      pinyin2ix=pinyin2ix,
                      pos2ix=pos2ix,
                      ix2glyph=ix2glyph,
                        device=device)
# -

' '.join(te_in[idx]),' '.join(predict), ' '.join(te_out[idx])

# +
# idxs = list(range(20))
idxs = [4,20,13,17]
preds = []
for i, idx in enumerate(idxs):
    data = te_in[idx]

    predict = beam_search_decode(model=model,
                                k=5,
                              bert=Anchibert,
                              tokenizer=tokenizer,
                              sent=data,
                              glyph2ix=glyph2ix,
                              pinyin2ix=pinyin2ix,
                              pos2ix=pos2ix,
                              ix2glyph=ix2glyph,
                                device=device)[0][0]
    preds.append(predict)
    
    

# -

for i, pred in enumerate(preds):
    print(f'sample {i}')
    print('top line: '+''.join(te_in[idxs[i]]))
    print('predict: '+ ''.join(pred))
    print('gold: ' + ''.join(te_out[idxs[i]]))
    print("================================")



StarcoderdataJupyterScriptsDedupFiltered