import streamlit as st
import torch
import numpy as np

tokenizer = torch.load('./tokenizer.pth')
bert_mlm_positive = torch.load('./bert_mlm_positive.pth', map_location=torch.device('cpu'))
bert_mlm_negative = torch.load('./bert_mlm_negative.pth', map_location=torch.device('cpu'))
bert_class = torch.load('./bert_classification.pth', map_location=torch.device('cpu'))
st.set_page_config(page_title="change the text style!", layout="centered")
st.markdown("## change the text style!")
phrase = st.text_input("phrase =", value='good!')
    
device = torch.device('cpu')    
def beam_translate(sent: str, num_iter: int = 4, M: int = 7):
    sentence_buf = [sent]
    for _ in range(num_iter):
        generate_var_buf = []
        for s in sentence_buf:
            gen_sentences = get_replacements(s, num_tokens=2, k_best=5)
            generate_var_buf += gen_sentences
        
        data_for_bert = tokenizer(generate_var_buf, padding=True, truncation=True, return_tensors="pt").to(device)
        p = bert_class(**data_for_bert)
        probs = p.logits.softmax(dim=-1)[0].detach().cpu().numpy()
        best_M = np.argsort(probs)[-M:]
        sentence_buf = []
        for i in best_M:
            sentence_buf.append(generate_var_buf[i])
    
    data_for_bert = tokenizer(generate_var_buf, padding=True, truncation=True, return_tensors="pt").to(device)
    p = bert_class(**data_for_bert)
    probs = p.logits.softmax(dim=-1)[0].detach().cpu().numpy()
    best = np.argsort(probs)[-1]
    return sentence_buf[best]
    
   
def get_replacements(sentence: str, num_tokens, k_best, epsilon=1e-3):
    """
    - split the sentence into tokens using the INGSOC-approved BERT tokenizer
    - find :num_tokens: tokens with the highest ratio (see above)
    - replace them with :k_best: words according to bert_mlm_positive
    :return: a list of all possible strings (up to k_best * num_tokens)
    """
    sentence_token = tokenizer(sentence, return_tensors='pt')
    sentence_token = {key: value.to(device) for key, value in sentence_token.items()}
    length = len(sentence_token['input_ids'][0])
    
    p_posit = bert_mlm_positive(**sentence_token).logits.softmax(dim=-1)[0]
    p_negat = bert_mlm_negative(**sentence_token).logits.softmax(dim=-1)[0]
    
    p_tokens_positive = p_posit[torch.arange(length), sentence_token['input_ids'][0]]
    p_tokens_negative = p_negat[torch.arange(length), sentence_token['input_ids'][0]]
    
    prob = ((p_tokens_positive + epsilon) / (p_tokens_negative + epsilon)).detach().cpu().numpy()
    replace_items = np.argsort(prob)[:num_tokens]
    results = []
    for i in replace_items:
        k_best_pos = np.argsort(p_posit[i].detach().cpu().numpy())[::-1][:k_best]
        for item in k_best_pos:
            new = sentence_token['input_ids'][0]
            new[i] = item
            results.append(' '.join(tokenizer.decode(new).split(' ')[1:-1]))
    return results
    
answer = beam_translate(phrase)
st.markdown(answer)