File size: 803 Bytes
86c445d 924d432 86c445d 924d432 86c445d 924d432 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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()
|