import tensorflow as tf 
import tensorflow_text as text
from tensorflow.train import Checkpoint
import pandas as pd
import numpy as np
import gradio as gr
from Model import Transformer

vocab = []
with open("vocab.txt", mode = "r", encoding = "utf-8") as file:
    for token in file:
        vocab.append(token.replace("\n", ""))
        
reserved_tokens=["[START]", "[END]", "[PAD]", "[UNK]"]
        
START = tf.argmax(tf.constant(reserved_tokens) == "[START]")
END = tf.argmax(tf.constant(reserved_tokens) == "[END]")
PAD = tf.argmax(tf.constant(reserved_tokens) == "[PAD]")        
VOCAB_SIZE = len(vocab)
print(VOCAB_SIZE)
D_MODEL = 256
NB_LAYERS = 6
FFN_UNITS = 2048
NB_PROJ = 8
DROPOUT_RATE = 0.1
MAX_LENGTH = 50

tokenizer = text.FastBertTokenizer(vocab, support_detokenization = True)
trimer = text.WaterfallTrimmer(max_seq_length = MAX_LENGTH)

transformer = Transformer(vocab_size_enc = VOCAB_SIZE, 
                          vocab_size_dec = 1,
                          d_model = D_MODEL,
                          nb_layers = NB_LAYERS,
                          FFN_units = FFN_UNITS,
                          nb_proj = NB_PROJ,
                          dropout_rate = DROPOUT_RATE)
                          
ckpt = Checkpoint()   
ckpt.restore("ckpt-10")  
print("Checkpoint Restaurado")  

def evaluate(sentence):
    ragged = tokenizer.tokenize([sentence])
    ragged = trimer.trim([ragged])[0]
    count = ragged.bounding_shape()[0]
    starts = tf.fill([count,1], START)
    ends = tf.fill([count,1], END)
    inputs = tf.concat([starts, ragged, ends], axis=1)
    inputs, _ = text.pad_model_inputs(inputs, max_seq_length = MAX_LENGTH + 2, pad_value = PAD)
    
    prediction = transformer(inputs, False)
    
    prediction = tf.round(prediction)
    print(prediction)
    if prediction == 0:
        return"Negative"
    else:
        return"Positive"
        
app = gr.Interface(fn = evaluate, title = "IMDb Sentiment Classifier", description = "Write a sentence with a positive or negative sentiment", inputs = "text", outputs = "text")   
app.launch(share = True)