import gradio as gr
from transformers import AutoFeatureExtractor, AutoTokenizer, SpeechEncoderDecoderModel
import torch

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)
    return classifier(sequence, labels, hypothesis_template=hypothesis_template, multi_class=True)

    


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="text")
iface.launch()