import gradio as gr
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
import logging
from typing import Tuple, Dict
import warnings
warnings.filterwarnings('ignore')

# Configuração de logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class QAModel:
    def __init__(self):
        logger.info("Inicializando modelo...")
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model_name = "pierreguillou/bert-base-cased-squad-v1.1-portuguese"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForQuestionAnswering.from_pretrained(self.model_name)
        self.model.to(self.device)
        logger.info(f"Modelo carregado no dispositivo: {self.device}")

    def answer_question(self, question: str, context: str) -> Tuple[str, Dict]:
        # Tokenização
        inputs = self.tokenizer.encode_plus(
            question, context,
            add_special_tokens=True,
            return_tensors="pt",
            max_length=512,
            truncation=True,
            padding="max_length"
        ).to(self.device)

        # Previsão
        with torch.no_grad():
            outputs = self.model(**inputs)

        # Processamento das respostas
        answer_start = torch.argmax(outputs.start_logits)
        answer_end = torch.argmax(outputs.end_logits)

        # Convertendo para tokens e depois para texto
        tokens = self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
        answer = tokens[answer_start:answer_end + 1]
        answer = self.tokenizer.convert_tokens_to_string(answer)

        # Calculando score
        confidence_start = torch.softmax(outputs.start_logits, dim=1)[0][answer_start].item()
        confidence_end = torch.softmax(outputs.end_logits, dim=1)[0][answer_end].item()
        confidence = (confidence_start + confidence_end) / 2

        return answer, {
            "confiança": f"{confidence:.2%}",
            "início": answer_start.item(),
            "fim": answer_end.item()
        }

def create_gradio_interface():
    # Instanciando o modelo
    qa_model = QAModel()

    def process_question(context: str, question: str) -> Tuple[str, str]:
        try:
            answer, metadata = qa_model.answer_question(question, context)
            return (
                answer,
                f"Confiança: {metadata['confiança']}\n"
                f"Posição da resposta: {metadata['início']} até {metadata['fim']}"
            )
        except Exception as e:
            logger.error(f"Erro ao processar pergunta: {str(e)}")
            return "Ocorreu um erro ao processar sua pergunta.", "N/A"

    # Criando a interface
    demo = gr.Interface(
        fn=process_question,
        inputs=[
            gr.Textbox(
                label="Contexto",
                placeholder="Cole aqui o texto de contexto...",
                lines=10
            ),
            gr.Textbox(
                label="Pergunta",
                placeholder="Digite sua pergunta sobre o texto acima...",
                lines=2
            )
        ],
        outputs=[
            gr.Textbox(label="Resposta"),
            gr.Textbox(label="Metadados")
        ],
        title="Perguntas e Respostas em Português",
        description="""
        Este aplicativo usa o modelo BERT em português para responder perguntas baseadas em um contexto fornecido.
        Cole um texto no campo de contexto e faça uma pergunta sobre ele.
        """,
        examples=[
            [
                "O Brasil é o maior país da América do Sul. Sua capital é Brasília e sua maior cidade é São Paulo. O país tem uma população de mais de 200 milhões de habitantes.",
                "Qual é a capital do Brasil?"
            ],
            [
                "A linguagem Python foi criada por Guido van Rossum em 1991. É uma linguagem de programação de alto nível, interpretada e orientada a objetos.",
                "Quem criou a linguagem Python?"
            ]
        ],
        theme=gr.themes.Soft()
    )

    return demo

if __name__ == "__main__":
    logger.info("Iniciando aplicação...")
    demo = create_gradio_interface()
    demo.launch(
        server_name="0.0.0.0",
        share=True,
        show_error=True
    )