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):
    hypothesis_template = 'Dette eksempelet er {}.'
    label_clean = str(labels).split(",")
    response = classifier(sequence, label_clean, hypothesis_template=hypothesis_template, multi_class=True)
    labels = response['labels']
    scores = response['scores']
    clean_output = {labels[idx]: float(scores[idx]) for idx in range(len(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.","Kutt smør i terninger, og la det temperere seg litt mens deigen elter. Ha hvetemel, sukker, gjær, salt og kardemomme i en bakebolle til kjøkkenmaskin. Bruker du fersk gjær kan du smuldre gjæren i bollen, eller røre den ut i melken. Alt vil ettehvert blande seg godt, så begge deler er like bra."]
example_labels=["politikk,helse,sport,religion", "helse,sport,religion, mat"]

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=10, 
        label="Possible candidate labels",
        placeholder="labels here...")],  
    outputs=gr.outputs.Label(num_top_classes=3),
    capture_session=True,
    interpretation="default"
    ,examples=[
        [example_text, example_labels]
    ])
iface.launch()