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)

    # 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

# Função para gerar várias notas e salvar em um arquivo
def gerar_varias_notas_e_salvar(feature_descriptions, arquivo_saida, num_notas_por_descricao=3):
    with open(arquivo_saida, 'w', encoding='utf-8') as file:
        for descricao in feature_descriptions:
            file.write(f"Prompt: {descricao}\n")
            file.write("Notas Geradas:\n")
            # Gerar várias notas para cada descrição
            notas_geradas = gerar_nota_release(descricao, max_new_tokens=50, num_return_sequences=num_notas_por_descricao)
            for i, nota in enumerate(notas_geradas):
                file.write(f"  Nota {i + 1}:\n")
                file.write(nota + "\n\n")

if __name__ == '__main__':
    # Lista de descrições de funcionalidades para gerar as notas
    feature_descriptions = [
        "Agora você pode verificar o quanto tem disponível em conta corrente para investir.",
        "Melhoramos a performance do app para navegação mais rápida.",
        "Agora você pode acessar relatórios detalhados sobre seus investimentos.",
        "Adicionamos uma nova seção de ajuda para facilitar o suporte aos usuários."
    ]

    # Gerar as notas de release e salvar no arquivo 'notas_de_release.txt'
    gerar_varias_notas_e_salvar(feature_descriptions, 'notas_de_release.txt', num_notas_por_descricao=3)

    print("Notas de release geradas e salvas em 'notas_de_release.txt'.")