from transformers import pipeline
import numpy as np
import gradio as gr

def netScores(tagList: list, sequence_to_classify: str, modelName: str) -> dict:
    classifier = pipeline("zero-shot-classification", model=modelName)
    hypothesis_template_pos = "This example is {}"
    hypothesis_template_neg = "This example is not {}"
    output_pos = classifier(sequence_to_classify, tagList, hypothesis_template=hypothesis_template_pos, multi_label=True)
    output_neg = classifier(sequence_to_classify, tagList, hypothesis_template=hypothesis_template_neg, multi_label=True)

    positive_scores = {}
    for x in range(len(tagList)):
        positive_scores[output_pos["labels"][x]] = output_pos["scores"][x]

    negative_scores = {}
    for x in range(len(tagList)):
        negative_scores[output_neg["labels"][x]] = output_neg["scores"][x]

    pos_neg_scores = {}
    for tag in tagList:
        pos_neg_scores[tag] = [positive_scores[tag],negative_scores[tag]]
    
    net_scores = {}
    for tag in tagList:
        net_scores[tag] = positive_scores[tag]-negative_scores[tag]

    net_scores = dict(sorted(net_scores.items(), key=lambda x:x[1], reverse=True))

    return net_scores

def compareTextAndLabels (userText, userLabels):
    userLabelsArray = userLabels.split(",")
    
    labelsScores = netScores (userLabelsArray, userText, 'akhtet/mDeBERTa-v3-base-myXNLI')
    for label in labelsScores:
        labelsScores[label] = str(np.round(labelsScores[label]*100,2))+"%"
    
    return labelsScores

               
demo = gr.Interface(
    fn=compareTextAndLabels,
    inputs=[gr.Textbox(label="Text"), gr.Textbox(label="Tags (separated by commas)")],
    outputs=[gr.Textbox(label="Tag Scores")],
)
demo.launch()