import gradio as gr
from transformers import pipeline
classifier = pipeline("zero-shot-classification", model="NbAiLab/nb-bert-base-mnli")


def sequence_to_classify(sequence, labels):
    response = { 'sequence': "Folkehelseinstituttets mest optimistiske anslag er at alle voksne er ferdigvaksinert innen midten av september.",'labels': ['helse', 'politikk', 'religion', 'sport'],'scores': [0.7680550217628479,0.21670468151569366,0.01563994586467743,0.00441053556278348]}
    clean_output = {idx: float(response['scores'].pop()) for idx in response['labels']}

    hypothesis_template = 'Dette eksempelet er {}.'
    label_clean = str(labels).split(",")
    response1 = classifier(sequence, label_clean, hypothesis_template=hypothesis_template, multi_class=True)
    labels = response1['labels']
    scores = response1['scores']
    clean_output1 = {idx: float(scores.pop(0)) for idx in label_clean}
    
    print("response is:{}".format(response))
    print(type(response))
    print("clean_output: {}".format(clean_output))
    print("\n")
    print("\n")
    print("response1 is:{}".format(response1))
    print(type(response1))
    print("clean_output1: {}".format(clean_output1))
    return clean_output1

example_text="Folkehelseinstituttets mest optimistiske anslag er at alle voksne er ferdigvaksinert innen midten av september."
example_labels=[["politikk", "helse", "sport", "religion"]]

def greet(name):
    return "Hello " + name + "!!"

iface = gr.Interface(
    title = "Zero-shot Classification of Norwegian Text",
    description = "Demo of zero-shot classification using NB-Bert base model (Norwegian).",
    fn=sequence_to_classify, 
    inputs=[gr.inputs.Textbox(lines=2, 
        label="Write a norwegian text you would like to classify...",
        placeholder="Text here..."),
        gr.inputs.Textbox(lines=2, 
        label="Possible candidate labels",
        placeholder="labels here...")],  
    outputs=gr.outputs.Label(num_top_classes=3, label="Categories"),
    capture_session=True,
    interpretation="default"
    ,examples=[
        [example_text, example_labels]
    ])
iface.launch()