File size: 2,592 Bytes
0bd62e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
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)
|