import gradio as gr
from transformers import AutoTokenizer

from nfqa_model import RobertaNFQAClassification


index_to_label = {0: 'NOT-A-QUESTION',
                  1: 'FACTOID',
                  2: 'DEBATE',
                  3: 'EVIDENCE-BASED',
                  4: 'INSTRUCTION',
                  5: 'REASON',
                  6: 'EXPERIENCE',
                  7: 'COMPARISON'}


model = RobertaNFQAClassification.from_pretrained("Lurunchik/nf-cats")
nfqa_tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")


def get_nfqa_prediction(text):
    output = model(**nfqa_tokenizer(text, return_tensors="pt"))
    index = output.logits.argmax()
    return index_to_label[int(index)]


iface = gr.Interface(fn=get_nfqa_prediction, inputs="text", outputs="text")
iface.launch()