import pandas as pd
import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel

# Carregando o modelo e o tokenizador do GPT-2
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

# Dados iniciais
data = {
    'Nome': ['Alice', 'Bob', 'Charlie'],
    'Idade': [25, 30, 35],
    'Cidade': ['Nova York', 'Los Angeles', 'Chicago'],
    'Feedback': [None, None, None]
}
df = pd.DataFrame(data)

# Função para responder perguntas com GPT-2
def answer_question_with_gpt(question):
    # Supondo que você queira incorporar dados do DataFrame na pergunta
    prompt = f"Considerando os dados: {df.to_string(index=False)}. Pergunta: {question} Resposta:"
    input_ids = tokenizer.encode(prompt, return_tensors='pt')
    max_length = len(input_ids[0]) + 50  # Define um limite máximo razoável para o comprimento da resposta
    generated_ids = model.generate(input_ids, max_length=max_length)
    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return generated_text.split("Resposta:")[1] if "Resposta:" in generated_text else generated_text

# Função para adicionar feedback
def add_feedback(nome, feedback):
    global df
    if nome in df['Nome'].values:
        df.loc[df['Nome'] == nome, 'Feedback'] = feedback
        return "Feedback adicionado com sucesso."
    else:
        return "Nome não encontrado no DataFrame."

with gr.Blocks() as demo:
    gr.Markdown("# Sistema de Consulta e Feedback de Dados")

    with gr.Row():
        with gr.Column():
            question_input = gr.Textbox(label="Faça uma Pergunta")
            answer_output = gr.Textbox(label="Resposta", interactive=False)
            ask_button = gr.Button("Perguntar")

        with gr.Column():
            name_input = gr.Textbox(label="Nome para Feedback")
            feedback_input = gr.Textbox(label="Feedback")
            feedback_result = gr.Textbox(label="Resultado do Feedback", interactive=False)
            submit_button = gr.Button("Enviar Feedback")

    ask_button.click(fn=answer_question_with_gpt, inputs=question_input, outputs=answer_output)
    submit_button.click(fn=add_feedback, inputs=[name_input, feedback_input], outputs=feedback_result)

demo.launch()