import gradio as gr
from transformers import pipeline
import concurrent.futures
import ktrain
import time

zero_shot = pipeline("zero-shot-classification")
distilbert = ktrain.load_predictor("models/distilbert-base-uncased-finetuned-internet-provider")
distilbert_v2 = ktrain.load_predictor("models/distilbert-base-uncased-finetuned-internet-provider-v2")

def zero_shot_predict(text):
    labels = ["Slow Connection", "Billing", "Setup", "No Connectivity"]
    preds = zero_shot(text, candidate_labels=labels)
    return {label: float(pred) for label, pred in zip(preds["labels"], preds["scores"])}

def distilbert_predict(text):
    labels = distilbert.get_classes()
    preds = distilbert.predict_proba(text)
    return {label: float(pred) for label, pred in zip(labels, preds)}

def distilbert_v2_predict(text):
    labels = distilbert_v2.get_classes()
    preds = distilbert_v2.predict_proba(text)
    return {label: float(pred) for label, pred in zip(labels, preds)}

def predict(text):
    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
        zero_shot_future = executor.submit(zero_shot_predict, text)
        distilbert_future = executor.submit(distilbert_predict, text)
        distilbert_v2_future = executor.submit(distilbert_v2_predict, text)
    concurrent.futures.wait([zero_shot_future, distilbert_future, distilbert_v2_future])
    zero_shot_preds = zero_shot_future.result()
    distilbert_preds = distilbert_future.result()
    distilbert_v2_preds = distilbert_v2_future.result()
    return zero_shot_preds, distilbert_preds, distilbert_v2_preds

input = gr.inputs.Textbox(label="Customer Sentence")
outputs = [gr.outputs.Label(num_top_classes=4, label="Zero-Shot-Classification"), gr.outputs.Label(num_top_classes=4, label="DistilBERT"), gr.outputs.Label(num_top_classes=4, label="DistilBERT v2")]
title = "Case Classification"
description = "Comparison of Zero-Shot-Classification and a fine-tuned DistilBERT."
gr.Interface(predict, input, outputs, live=False, live_update=False, title=title, analytics_enabled=False,
             description=description, capture_session=True).launch()