Spaces:
Runtime error
Runtime error
import torch | |
import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForQuestionAnswering | |
tokenizer = AutoTokenizer.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad") | |
model = AutoModelForQuestionAnswering.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad") | |
def QA_function(context, question): | |
inputs = tokenizer(question, context, add_special_tokens=True, return_tensors="pt") | |
input_ids = inputs["input_ids"].tolist()[0] | |
outputs = model(**inputs) | |
answer_start_scores = outputs.start_logits | |
answer_end_scores = outputs.end_logits | |
# Get the most likely beginning of answer with the argmax of the score | |
answer_start = torch.argmax(answer_start_scores) | |
# Get the most likely end of answer with the argmax of the score | |
answer_end = torch.argmax(answer_end_scores) + 1 | |
answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[answer_start:answer_end])) | |
return answer | |
gradio_ui = gr.Interface(QA_function, [gr.inputs.Textbox(lines=7, label="Context"), gr.inputs.Textbox(label="Question")], gr.outputs.Textbox(label="Answer")) | |
gradio_ui.launch() |