import streamlit as st
from transformers import T5Tokenizer, T5ForConditionalGeneration, pipeline

# Substitua 'seu-nome-de-usuário/nome-do-modelo' pelo caminho correto do seu modelo no Hugging Face Hub
MODEL_NAME = "exo-is/t5-small-60M-esg-keyword"

@st.cache_resource
def load_model():
    tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
    model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME)
    return pipeline('text2text-generation', model=model, tokenizer=tokenizer)

st.title('Gerador de Texto T5')

gerador = load_model()

# Entrada do usuário
prompt = st.text_area('Digite o prompt inicial:')

# Configurações de geração
max_length = st.slider('Comprimento máximo do texto gerado:', min_value=10, max_value=500, value=100)
num_return_sequences = st.number_input('Número de sequências a gerar:', min_value=1, max_value=5, value=1)

if st.button('Gerar Texto'):
    if prompt:
        with st.spinner('Gerando texto...'):
            resultados = gerador(prompt, max_length=max_length, num_return_sequences=num_return_sequences)
        
        st.write('Texto(s) gerado(s):')
        for i, resultado in enumerate(resultados, 1):
            st.write(f"Sequência {i}:")
            st.write(resultado['generated_text'])
    else:
        st.warning('Por favor, digite um prompt inicial.')