import gradio as gr
import datetime
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)
    predicted_labels = response['labels']
    predicted_scores = response['scores']
    clean_output = {idx: float(predicted_scores.pop(0)) for idx in predicted_labels}
    print("Date:{} , Sequece:{}, Labels: {}".format(
        str(datetime.datetime.now()),
        sequence,
        predicted_labels)
    )
    return clean_output

example_text1="Folkehelseinstituttets mest optimistiske anslag er at alle voksne er ferdigvaksinert innen midten av september."
example_labels1="politikk,helse,sport,religion"
example_text2="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_labels2="helse,sport,religion, mat"

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_text1, example_labels1],
        [example_text2, example_labels2]
    ])
iface.launch()