from transformers import DistilBertForQuestionAnswering, DistilBertConfig, DistilBertTokenizerFast
import torch

model = DistilBertForQuestionAnswering(DistilBertConfig.from_pretrained('distilbert/distilbert-base-multilingual-cased')).to("cpu")
st_dict = torch.load("QazDistilBERT.pt",map_location=torch.device('cpu'))
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

example_contexts = [
    "Қазақстанның ұлттық құрамы алуан түрлі. Халықтың басым бөлігін тұрғылықты қазақ халқы құрайды, пайыздық үлесі — 70,18%[10], орыстар — 18,42%, өзбектер — 3,29%, украиндар — 1,36%, ұйғырлар — 1,48%, татарлар — 1,06%, басқа халықтар 5,38%.[11] Халықтың 75% астамын мұсылмандар құрайды, православты христиандар — 21%, қалғаны басқа да дін өкілдері.[12]",
    "Қазақстан бес мемлекетпен шекаралас, соның ішінде әлемдегі құрлықтағы ең ұзын шекара, солтүстігінде және батысында Ресеймен — 7591 км құрайды. Оңтүстігінде: Түрікменстан — 426 км, Өзбекстан — 2354 км және Қырғызстан — 1241 км, ал шығысында: Қытаймен — 1782 км шектеседі. Жалпы құрлық шекарасының ұзындығы — 13394 км. Батыста Каспий көлімен (2000 км), оңтүстік батыста Арал теңізімен шайылады.[9] 2024 жылдың 1 наурыздағы елдегі тұрғындар саны — 20 075 271[4], бұл әлем бойынша 64-орын. Жер көлемі жағынан әлем елдерінің ішінде 9-орын алады (2 724 902 км²).",
    "Қазақстан — 1995 жылғы 30 тамыздағы республикалық референдумда қабылданған Конституция бойынша — өзін демократиялы, зайырлы, құқықты және әлеуметті мемлекет ретінде орнықтырды. Қазақстан Республикасы – президенттік басқару формасындағы біртұтас мемлекет. Республиканың ең жоғарғы өкілді органы — Парламент. Ол республиканың заң шығару құзіретін жүзеге асырады."
]
example_questions = [
"Қазақстанның халқы неше пайызды қазақтар құрайды?",
"Қазақстан нешеу мемлекетпен шекаралас?",
"Қазақстандағы басқару формасы қандай?",
]

examples = [[context, question] for context, question in zip(example_contexts, example_questions)]

# Создаем интерфейс
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="Введите контекст и задайте вопрос, чтобы получить ответ.",
    examples=examples
)

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