import gradio as gr
import os
import torch

from model import create_roberta_model
from timeit import default_timer as timer
from typing import Tuple, Dict

# Setup class names
with open("class_names.txt", "r") as f:
  class_names = [name.strip() for name in f.readlines()]

### Load example texts ###
example_texts = []
with open("example_texts.txt", "r") as file:
    example_texts = [line.strip() for line in file.readlines()]

### Model and transforms preparation ###
# Create model and tokenizer
model, tokenizer = create_roberta_model(output_shape=len(class_names))

# Load saved weights
model.load_state_dict(
    torch.load(f="roberta-base.pth",
               map_location=torch.device("cpu")) # load to CPU
)

### Predict function ###
def predict(text) -> Tuple[Dict, float]:
  # Start a timer
  start_time = timer()

  # Set the model to eval
  model.eval()

  # Set up the inputs
  X = tokenizer(text, padding="max_length", truncation=True, return_tensors='pt')

  # Put model into eval mode, make prediction
  model.eval()
  with torch.inference_mode():
    # Pass tokenized text through the model and turn the prediction logits into probaiblities
    pred_probs = torch.softmax(model(**X).logits, dim=1)

  # Create a prediction label and prediction probability dictionary
  pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}

  # Calculate pred time
  end_time = timer()
  pred_time = round(end_time - start_time, 4)

  # Return pred dict and pred time
  return pred_labels_and_probs, pred_time

### 4. Gradio app ###
# Create title, description and article
title = "A roberta-base Classifier"
description = "[A roberta-base BERT based model](https://huggingface.co/roberta-base) text model to classify text on the [HuggingFace 🤗 dair-ai/emotion dataset](https://huggingface.co/datasets/dair-ai/emotion). [Source Code Found Here](https://colab.research.google.com/drive/1P7rfiDF1jfNHKmkB7WjHPi8PQBLQ4Ege?usp=sharing)"
article = "Built with [Gradio](https://github.com/gradio-app/gradio) and [PyTorch](https://pytorch.org/). [Source Code Found Here](https://colab.research.google.com/drive/1P7rfiDF1jfNHKmkB7WjHPi8PQBLQ4Ege?usp=sharing)"

# Create the Gradio demo
demo = gr.Interface(fn=predict,
                    inputs=gr.Textbox(lines=2, placeholder="Type your text here..."),
                    outputs=[gr.Label(num_top_classes=5, label="Predictions"),
                             gr.Number(label="Prediction time (s)")],
                    examples=example_texts,
                    title=title,
                    description=description,
                    article=article)

# Launch the demo
demo.launch()