bot_innopolis / app.py
cody82's picture
Update app.py
2e03cda verified
raw
history blame
1.32 kB
import gradio as gr
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
model_name = "deepset/roberta-base-squad2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
context = (
"Университет Иннополис был основан в 2012 году. "
"Это современный вуз в России, специализирующийся на IT и робототехнике, "
"расположенный в городе Иннополис, Татарстан."
)
def respond(question, history=None):
if history is None:
history = []
inputs = tokenizer.encode_plus(question, context, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model(**inputs)
start_scores = outputs.start_logits
end_scores = outputs.end_logits
start = torch.argmax(start_scores)
end = torch.argmax(end_scores) + 1
answer_tokens = inputs['input_ids'][0][start:end]
answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
history.append((question, answer))
return history
iface = gr.ChatInterface(fn=respond, title="Innopolis Q&A")
iface.launch()