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-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 predict(text):
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        zero_shot_future = executor.submit(zero_shot_predict, text)
        distilbert_future = executor.submit(distilbert_predict, text)
    concurrent.futures.wait([zero_shot_future, distilbert_future])
    zero_shot_preds = zero_shot_future.result()
    distilbert_preds = distilbert_future.result()
    return zero_shot_preds, distilbert_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")]
title = "Case Template 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()