|
|
from transformers import AutoTokenizer, AutoModelForQuestionAnswering |
|
|
import gradio as gr |
|
|
import torch, torchvision |
|
|
|
|
|
|
|
|
model_name = "mbwolff/distilbert-base-uncased-finetuned-squad" |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
model = AutoModelForQuestionAnswering.from_pretrained(model_name) |
|
|
|
|
|
def answer_question(question, context): |
|
|
""" |
|
|
Answers a question based on a given context. |
|
|
""" |
|
|
inputs = tokenizer.encode_plus(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 |
|
|
answer_start = torch.argmax(answer_start_scores) |
|
|
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 |
|
|
|
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=answer_question, |
|
|
inputs=[ |
|
|
|
|
|
|
|
|
gr.Textbox(label="Enter your question here...", lines=2), |
|
|
gr.Textbox(label="Enter the context here...", lines=5) |
|
|
], |
|
|
outputs="text", |
|
|
title="Question Answering Chatbot", |
|
|
description="Ask a question and provide a context, and the chatbot will try to answer it." |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
iface.launch() |