import gradio as gr
from gradio.mix import Parallel
from transformers import pipeline
import torch

question = "Ze erakundea immatrikulatu ditu ondasunak?"
context = "Bilbo, Donostia eta Gasteizko apezpikutzen arabera, elizbarrutiek ez diote Estatuari immatrikulatutako ondasunik itzuli behar - erregistraturik ez zeuden eta aurrenekoz Eliza Katolikoak inskribatu zituen ondasunak-. Hala ere, Gasteizko elizbarrutiak uste du erregularizatutako zortzi ondasunen inguruan zalantzak daudela.\n1998tik 2015era Elizak immatrikulatutako 1.000 bat ondasunen zerrenda zabaldu zuen atzo Gotzainen Biltzarrak. Arrazoi ezberdinengatik jabetza egiaztatzerik izan ez eta Estatuari bueltatzeko aukera onartu du Espainiako Elizak, Gobernuaren eta erakunde erlijiosoaren artean adostutako akordioaren arabera.\nHipoteka legearen 206. artikuluko bide berezia baliatuz 17 urte horietan ia 35.000 ondasun immatrikulatu zituzten elizbarrutiek. Beraz, zerrendan agertzen direnak % 3 baino gutxiago dira."

qa1 = pipeline("question-answering", model="MarcBrun/ixambert-finetuned-squad")
qa2 = pipeline("question-answering", model="MarcBrun/ixambert-finetuned-squad-eu")
qa3 = pipeline("question-answering", model="MarcBrun/ixambert-finetuned-squad-eu-en")

def answer_1(question, context):
  answer = qa1(question=question, context=context)
  return answer["answer"]
 
def answer_2(question, context):
  answer = qa2(question=question, context=context)
  return answer["answer"]
 
def answer_3(question, context):
  answer = qa3(question=question, context=context)
  return answer["answer"]

iface1 = gr.Interface(fn=answer_1, inputs=[gr.inputs.Textbox(),gr.inputs.Textbox()], outputs=gr.outputs.Textbox(label="ixambert-squad"))

iface2 = gr.Interface(fn=answer_2, inputs=[gr.inputs.Textbox(),gr.inputs.Textbox()], outputs=gr.outputs.Textbox(label="ixambert-squad-eu"))

iface3 = gr.Interface(fn=answer_3, inputs=[gr.inputs.Textbox(),gr.inputs.Textbox()], outputs=gr.outputs.Textbox(label="ixambert-squad-eu-en"))

Parallel(iface1, iface2, iface3, inputs=[gr.inputs.Textbox(lines=1, default=question, placeholder="Question Here...", label="Question"),gr.inputs.Textbox(lines=7, default=context, placeholder="Context Here...", label="Context")], title = "Question Answering in Basque", description="This demo compares the outputs of three QA models based on IXAmBERT. Each has been trained on different data: one on SQuADv1.1, another on an experimental version of SQuAD in Basque, and the last on both those datasets.").launch()