from tamilatis.predict import TamilATISPredictor 
from tamilatis.model import JointATISModel
import numpy as np
from sklearn.preprocessing import LabelEncoder
import gradio as gr



model_name = "microsoft/xlm-align-base"
tokenizer_name = "microsoft/xlm-align-base"
num_labels = 78
num_intents = 23
checkpoint_path = "models/xlm_align_base.bin"
intent_encoder_path = "models/intent_classes.npy"
ner_encoder_path = "models/ner_classes.npy"

def predict_function(text):
  label_encoder = LabelEncoder()
  label_encoder.classes_ = np.load(ner_encoder_path)
  
  intent_encoder = LabelEncoder()
  intent_encoder.classes_ = np.load(intent_encoder_path)
  
  model  = JointATISModel(model_name,num_labels,num_intents)
  predictor = TamilATISPredictor(model,checkpoint_path,tokenizer_name,
                              label_encoder,intent_encoder,num_labels)
  slot_prediction, intent_preds = predictor.get_predictions(text)
  return slot_prediction, intent_preds


title = "MultiTask Learning in Intent Detection and Slot Prediction for Tamil Conversational Dialogues using Multilingual Pretrained Models"  
article="This is a demo for the MultiTask model trained on Tamil Translated conversations from ATIS dataset. The code can be found [here](https://github.com/seanbenhur/tamilatis). Made with ❤ by [Sean Benhur](https://www.linkedin.com/in/seanbenhur/)" 
examples = ["ஹைதராபாத்தில் இருந்து உதய்பூர் செல்லும் விமானங்களைக் காட்டு", "எனக்கு டெல்லியில் இருந்து சென்னைக்கு விமானம் வேண்டும்"]

intent_output = gr.outputs.Textbox(type="auto",label="Intent")
slots_output = gr.outputs.Textbox(type="auto",label="Slots")
iface = gr.Interface(fn=predict_function,article=article, inputs="text", title=title,outputs=[intent_output,slots_output],
examples=examples)
iface.launch()