import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.datasets import imdb  # pyright: reportMissingImports=false
from huggingface_hub import from_pretrained_keras
import gradio as gr
from typing import Dict

class KerasIMDBTokenizer:

    def __init__(self, vocab_size: int = 20000) -> None:

        # Parameters used in `keras.datasets.imdb.load_data`
        self.START_CHAR = 1
        self.OOV_CHAR = 2
        self.INDEX_FROM = 3

        self.word_index: dict[str, int] = imdb.get_word_index()
        self.word_index = {
            token: input_id + self.INDEX_FROM
            for token, input_id in self.word_index.items() if input_id <= vocab_size
        }

    def tokenize_and_pad(self, text: str, maxlen: int = 200) -> np.ndarray:
        tokens = text.split()
        input_ids = [self.word_index.get(token.lower(), self.OOV_CHAR) for token in tokens]
        input_ids.insert(0, self.START_CHAR)
        # pad_sequences only accepts a list of sequences
        return pad_sequences([input_ids], maxlen=maxlen)

model = from_pretrained_keras("keras-io/text-classification-with-transformer", compile=False)
tokenizer = KerasIMDBTokenizer()

def sentiment_analysis(model_input: str) -> Dict[str, float]:
    tokenized = tokenizer.tokenize_and_pad(model_input)

    prediction = model.predict(tokenized)[0]
    ret = {
        "negative": float(prediction[0]),
        "positive": float(prediction[1])
    }
    return ret


model_input = gr.Textbox("Input text here", show_label=False)
model_output = gr.Label("Sentiment Analysis Result", num_top_classes=2, show_label=True, label="Sentiment Analysis Result")


examples = [
    (
        "Story of a man who has unnatural feelings for a pig. "
        "Starts out with a opening scene that is a terrific example of absurd comedy. "
        "A formal orchestra audience is turned into an insane, violent mob by the crazy chantings of it's singers. "
        "Unfortunately it stays absurd the WHOLE time with no general narrative eventually making it just too off putting. "
        "Even those from the era should be turned off. "
        "The cryptic dialogue would make Shakespeare seem easy to a third grader. "
        "On a technical level it's better than you might think with some good cinematography by future great Vilmos Zsigmond. "
        "Future stars Sally Kirkland and Frederic Forrest can be seen briefly."
    ),
    (
        "I came in in the middle of this film so I had no idea about any credits or even its title till I looked it up here, "
        "where I see that it has received a mixed reception by your commentators. "
        "I'm on the positive side regarding this film but one thing really caught my attention as I watched: "
        "the beautiful and sensitive score written in a Coplandesque Americana style. "
        "My surprise was great when I discovered the score to have been written by none other than John Williams himself. "
        "True he has written sensitive and poignant scores such as Schindler's List but one usually associates "
        "his name with such bombasticities as Star Wars. "
        "But in my opinion what Williams has written for this movie surpasses anything I've ever heard of his "
        "for tenderness, sensitivity and beauty, fully in keeping with the tender and lovely plot of the movie. "
        "And another recent score of his, for Catch Me if You Can, shows still more wit and sophistication. "
        "As to Stanley and Iris, I like education movies like How Green was my Valley and Konrack, "
        "that one with John Voigt and his young African American charges in South Carolina, "
        "and Danny deVito's Renaissance Man, etc. They tell a necessary story of intellectual and spiritual awakening, "
        "a story which can't be told often enough. This one is an excellent addition to that genre."
    )
]

title = "Text classification with Transformer"
description = "Implement a Transformer block as a Keras layer and use it for text classification."
article = (
    "Author: Xin Sui "
    "Based on <a href=\"https://keras.io/examples/nlp/text_classification_with_transformer\">this</a> "
    "keras example by <a href=\"https://twitter.com/NandanApoorv\">Apoorv Nandan</a>. "
    "HuggingFace Model <a href=\"https://huggingface.co/keras-io/text-classification-with-transformer\">here</a>"
)


app = gr.Interface(
    sentiment_analysis, 
	inputs=model_input,
    outputs=model_output, 
    examples=examples,
    title=title,
	description=description,
    article=article,
    allow_flagging='never',
    analytics_enabled=False,
)

app.launch(enable_queue=True)