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):
    #sequence_to_classify = 'Folkehelseinstituttets mest optimistiske anslag er at alle voksne er ferdigvaksinert innen midten av september.'
    #candidate_labels = ['politikk', 'helse', 'sport', 'religion']
    hypothesis_template = 'Dette eksempelet er {}.'
    #classifier(sequence_to_classify, candidate_labels, hypothesis_template=hypothesis_template, multi_class=True)
    response = classifier(sequence, labels, hypothesis_template=hypothesis_template, multi_class=True)
    clean_output = {labels[idx]: str(respose['scores'].pop(1)) for idx in response['labels']}
    print("response is:{}".format(response))
    print("clean_output: {}".format(clean_output))
    return clean_output

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()