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

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

processing = False

df = pd.read_csv('anomalies.csv')

def response(question, history):
    global processing
    processing = True
    
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    prompt = f"Considerando os dados: {df.to_string(index=False)}. Pergunta: {question} Resposta:"
    inputs = tokenizer(prompt, return_tensors='pt', padding='max_length', truncation=True, max_length=512)
    attention_mask = inputs['attention_mask']
    input_ids = inputs['input_ids']
    
    generated_ids = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_length=len(input_ids[0]) + 100,  # Aumentar o limite de geração
        temperature=0.65,  # Ajustar a criatividade
        top_p=0.9,  # Usar nucleus sampling
        no_repeat_ngram_size=2  # Evitar repetições desnecessárias
    )
    
    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    # Processando para extrair apenas a resposta após "Resposta:"
    response_part = generated_text.split("Resposta:")[1] if "Resposta:" in generated_text else "Resposta não encontrada."
    # Limpeza adicional para remover qualquer texto indesejado após a resposta
    final_response = response_part.split(".")[0] + "."  # Isso assume que a resposta termina na primeira sentença.
    
    history.append((question, final_response))
    processing = False
    return "", history

def loading():
    return "Carregando ..."


with gr.Blocks() as app:
    with gr.Column(elem_id="column_container"):
        chatbot = gr.Chatbot([], elem_id="chatbot")
        with gr.Column():
            send = gr.Label(value="Escreva sua QUESTÃO abaixo e pressione ENTER")
        query = gr.Textbox(
            label="Escreva sua questão aqui:",
            placeholder="Sobre o que você quer saber?",
        )
        clear = gr.ClearButton([query, chatbot])
    query.submit(response, [query, chatbot], [query, chatbot], queue=True)

app.launch()