from transformers import DistilBertForQuestionAnswering, DistilBertConfig, DistilBertTokenizerFast
import torch

model = DistilBertForQuestionAnswering(DistilBertConfig.from_pretrained('distilbert/distilbert-base-multilingual-cased')).to("cpu")
st_dict = torch.load("save/best_f1/checkpoint/QazDistilBERT.pt")
model.load_state_dict(st_dict)
tokenizer = DistilBertTokenizerFast.from_pretrained("dappyx/QazDistilbertFast-tokenizerV3")

import gradio as gr

def qa_pipeline(text,question):
  inputs = tokenizer(question, text, return_tensors="pt")
  input_ids = inputs['input_ids'].to("cpu")
  attention_mask = inputs['attention_mask'].to("cpu")
  outputs = model(input_ids=input_ids,attention_mask=attention_mask)

  start_index = torch.argmax(outputs.start_logits, dim=-1).item()
  end_index = torch.argmax(outputs.end_logits, dim=-1).item()

  predict_answer_tokens = inputs.input_ids[0, start_index : end_index + 1]
  return tokenizer.decode(predict_answer_tokens)

def answer_question(context, question):
    result = qa_pipeline(context, question)
    return result


# Создаем интерфейс
iface = gr.Interface(
    fn=answer_question, 
    inputs=[
        gr.Textbox(lines=10, label="Context"),
        gr.Textbox(lines=2, label="Question")
    ], 
    outputs="text",
    title="Question Answering Model",
    description="Введите контекст и задайте вопрос, чтобы получить ответ."
)

# Запускаем интерфейс
iface.launch()