import gradio as gr

# Load model from training checkpoint
from simpletransformers.question_answering import QuestionAnsweringModel, QuestionAnsweringArgs

model = QuestionAnsweringModel("bert", "./bert_outputs/bert/best_model", use_cuda=False)

# Função para realizar predições
def perform_prediction(context, question):
    to_predict = [{"context": context, "qas": [{"question": question, "id": "0"}]}]
    answers, probabilities = model.predict(to_predict, n_best_size=2)
    return answers[0]["answer"]

# Criando a interface do Gradio
iface = gr.Interface(
    fn=perform_prediction,
    inputs=[
        gr.Textbox(placeholder="Insira o contexto aqui...", label="Contexto"),
        gr.Textbox(placeholder="Faça uma pergunta sobre o contexto...", label="Pergunta")
    ],
    outputs=gr.Textbox(label="Resposta"),
    live=True,
    theme="compact",
)

# Iniciando a interface
iface.launch()