import torch
import gradio as gr
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline

model_eng   = "deepset/roberta-base-squad2"
model_br    = "mrm8488/bert-base-portuguese-cased-finetuned-squad-v1-pt"

PIPELINES = {
    "English": pipeline(
        "question-answering",
        model="deepset/roberta-base-squad2",
        tokenizer="deepset/roberta-base-squad2"
    ),
    "Português": pipeline(
        "question-answering",
        model="mrm8488/bert-base-portuguese-cased-finetuned-squad-v1-pt",
        tokenizer="mrm8488/bert-base-portuguese-cased-finetuned-squad-v1-pt"
    )
}

def read_file(file_content):
    try:
        content = file_content.read()
        return content
    except Exception as e:
        return f"Ocorreu um erro: {e}"
    
def get_answer(lang , text_content, question, file_content):
    if text_content and file_content:
        return "Selecione apenas uma opção!"
    
    if not text_content and not file_content:
        return "Preencha algum dos campos de conteudo antes para ter uma resposta."
    
    if file_content:
        content = file_content
    else:
        content = text_content

    pipe = PIPELINES.get(lang)
    if pipe is None:
        return "Idioma não suportado."
    
    res = pipe(
        context=content,
        question=question,
    )
    return res

demo = gr.Interface(fn=get_answer,
                    inputs=[
                    gr.Radio(["English", "Português"], label="Selecione o idioma"),
                    gr.Textbox(label="Contexto (texto livre)", placeholder="Digite seu texto aqui...", lines=6),
                    gr.Textbox(label="Pergunta", placeholder="Digite sua pergunta aqui...", lines=2),
                    gr.File(label="Ou faça upload de um arquivo .txt", file_types=['.txt'])
                    ],
                    outputs=gr.Textbox(label="Resposta"),
                    title="QnA Multilíngue",
                    description="Faça perguntas em inglês ou português. Carrega o modelo adequado automaticamente.",
                    )

if __name__ == "__main__":
    demo.launch()