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()