import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch

# Carregar o tokenizer e o modelo ajustado
tokenizer = GPT2Tokenizer.from_pretrained('modelo_treinado')
model = GPT2LMHeadModel.from_pretrained('modelo_treinado')

# Configurar o dispositivo para GPU se disponível
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Função para gerar notas de release
def gerar_nota_release(feature_description_prompt, max_new_tokens=50, num_return_sequences=1):
    # Definir o template com um placeholder para a descrição da funcionalidade
    template = """É hora de atualizar o seu íon Itaú.
    {}
Atualize o app já e aproveite!"""

    # Formatar o template com o placeholder para o prompt
    prompt = template.format(feature_description_prompt)

    # Tokenizar o prompt de entrada
    inputs = tokenizer.encode(prompt, return_tensors='pt').to(device)

    # Obter o comprimento do prompt
    prompt_length = inputs.shape[1]

    # Gerar texto
    outputs = model.generate(
        inputs,
        max_new_tokens=max_new_tokens,
        num_return_sequences=num_return_sequences,
        no_repeat_ngram_size=2,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.7,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )

    # Decodificar e retornar o texto gerado
    notas = []
    for i in range(num_return_sequences):
        # Decodificar a sequência gerada
        output = outputs[i]
        texto_completo = tokenizer.decode(output, skip_special_tokens=True)

        # Extrair apenas a parte gerada após o prompt
        texto_gerado = texto_completo[len(prompt):].strip()

        # Montar a nota de release completa
        nota_release = template.format(feature_description_prompt + ' ' + texto_gerado)
        notas.append(nota_release)
    return notas[0]

# Definir a função para o chat
def chat_model(feature_description):
    return gerar_nota_release(feature_description)

# Criar a interface do Gradio
chat_interface = gr.Interface(
    fn=chat_model,
    inputs=gr.Textbox(lines=2, placeholder="Digite a descrição da funcionalidade..."),
    outputs=gr.Textbox(),
    title="Chat para Geração de Notas de Release",
    description="Digite uma descrição de funcionalidade e o modelo gerará uma nota de release de acordo com o modelo treinado.",
)

# Executar a interface
if __name__ == "__main__":
    chat_interface.launch(share=True)