import streamlit as st

import torch
import torch.nn as nn
import torch.nn.functional as F
import transformers
import pandas as pd

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from transformers import MarianMTModel, MarianTokenizer
model_name = 'Helsinki-NLP/opus-mt-ROMANCE-en'

@st.cache_resource
def get_tokenizer(model_name):
    return MarianTokenizer.from_pretrained(model_name)

@st.cache_resource
def get_model(model_name):
    return MarianMTModel.from_pretrained(model_name).to(device)

tokenizer = get_tokenizer(model_name)
model = get_model(model_name)

print(f"The model has {model.num_parameters():,d} parameters.")

input_text = st.text_input("Enter text to translate", "Hola, mi nombre es Juan")
input_text = input_text.strip()
if not input_text:
    st.stop()

output_so_far = st.text_input("Enter text translated so far", "Hello, my")

input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)

# tokenize the output so far
with tokenizer.as_target_tokenizer():
    output_tokens = tokenizer.tokenize(output_so_far)
    decoder_input_ids = tokenizer.convert_tokens_to_ids(output_tokens)

# Add the start token
decoder_input_ids = [model.config.decoder_start_token_id] + decoder_input_ids

with torch.no_grad():
    model_output = model(
        input_ids = input_ids,
        decoder_input_ids = torch.tensor([decoder_input_ids]).to(device))


last_token_logits = model_output.logits[0, -1].cpu()
assert len(last_token_logits.shape) == 1
most_likely_tokens = last_token_logits.topk(k=5)

probs = last_token_logits.softmax(dim=-1)
probs_for_likely_tokens = probs[most_likely_tokens.indices]

with tokenizer.as_target_tokenizer():
    probs_table = pd.DataFrame({
        'token': [tokenizer.decode(token_id) for token_id in most_likely_tokens.indices],
        'id': most_likely_tokens.indices,
        'probability': probs_for_likely_tokens,
        'logprob': probs_for_likely_tokens.log(),
        'cumulative probability': probs_for_likely_tokens.cumsum(0)
    })


st.write(probs_table)
st.write(model.config.decoder_start_token_id)