import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import traceback

# Desative os avisos para uma saída mais limpa
import warnings
warnings.filterwarnings("ignore")

MODEL_NAME = 'exo-is/t5-small-60M-esg-keyword'

@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
    return model, tokenizer

st.title('Gerador de Texto T5')

try:
    model, tokenizer = load_model()
    gerador = pipeline('text2text-generation', model=model, tokenizer=tokenizer)

    st.write("Este aplicativo gera texto baseado em um prompt inicial usando o modelo T5.")
    
    prompt = st.text_area("Digite o prompt inicial:", height=100)
    
    col1, col2 = st.columns(2)
    with col1:
        max_length = st.slider("Comprimento máximo do texto gerado:", min_value=10, max_value=500, value=100)
    with col2:
        num_return_sequences = st.slider("Número de sequências a gerar:", min_value=1, max_value=5, value=1)
    
    temperature = st.slider("Temperatura (criatividade):", min_value=0.1, max_value=1.0, value=0.7, step=0.1)
    
    if st.button('Gerar Texto'):
        if prompt:
            with st.spinner('Gerando texto...'):
                # Usar beam search se num_return_sequences > 1
                if num_return_sequences > 1:
                    resultados = gerador(prompt, max_length=max_length, num_return_sequences=num_return_sequences, 
                                         temperature=temperature, num_beams=num_return_sequences, 
                                         no_repeat_ngram_size=2, early_stopping=True)
                else:
                    resultados = gerador(prompt, max_length=max_length, num_return_sequences=1, 
                                         temperature=temperature, do_sample=True)
            
            st.subheader('Texto(s) gerado(s):')
            for i, resultado in enumerate(resultados, 1):
                st.markdown(f"**Sequência {i}:**")
                st.write(resultado['generated_text'])
                st.markdown("---")
        else:
            st.warning('Por favor, digite um prompt inicial.')

except Exception as e:
    st.error(f"Um erro ocorreu: {str(e)}")
    st.error("Traceback completo:")
    st.error(traceback.format_exc())

st.markdown("---")
st.write("Desenvolvido com ❤️ usando Streamlit e Hugging Face Transformers")