BugZoid/text-humanizer

1import streamlit as st2from transformers import T5ForConditionalGeneration, T5Tokenizer3import torch4from torch.utils.data import Dataset, DataLoader5import json6import os7from datetime import datetime8import tempfile9 10# Custom dataset for fine-tuning11class TextHumanizerDataset(Dataset):12    def __init__(self, data, tokenizer, max_length=512):13        self.data = data14        self.tokenizer = tokenizer15        self.max_length = max_length16 17    def __len__(self):18        return len(self.data)19 20    def __getitem__(self, idx):21        item = self.data[idx]22        input_encoding = self.tokenizer(23            f"reescreva em português natural, mantendo todas as informações: {item['input_text']}",24            max_length=self.max_length,25            padding='max_length',26            truncation=True,27            return_tensors='pt'28        )29        30        target_encoding = self.tokenizer(31            item['output_text'],32            max_length=self.max_length,33            padding='max_length',34            truncation=True,35            return_tensors='pt'36        )37 38        return {39            'input_ids': input_encoding['input_ids'].squeeze(),40            'attention_mask': input_encoding['attention_mask'].squeeze(),41            'labels': target_encoding['input_ids'].squeeze()42        }43 44def get_storage_path():45    """Retorna o caminho correto para armazenamento no Hugging Face Spaces"""46    if os.environ.get('SPACE_ID'):  # Verifica se está rodando no Spaces47        return '/data'  # Diretório persistente no Spaces48    else:49        # Fallback para desenvolvimento local50        temp_dir = tempfile.gettempdir()51        feedback_dir = os.path.join(temp_dir, 'feedback_data')52        os.makedirs(feedback_dir, exist_ok=True)53        return feedback_dir54 55def save_feedback(input_text, output_text, rating):56    """Salva o feedback do usuário para futuro treinamento"""57    feedback_data = {58        'input_text': input_text,59        'output_text': output_text,60        'rating': rating,61        'timestamp': datetime.now().isoformat()62    }63    64    storage_path = get_storage_path()65    feedback_file = os.path.join(storage_path, 'feedback.json')66    67    try:68        # Cria arquivo se não existir69        if not os.path.exists(feedback_file):70            with open(feedback_file, 'w') as f:71                f.write('')72        73        # Append do novo feedback74        with open(feedback_file, 'a') as f:75            f.write(json.dumps(feedback_data) + '\n')76            77        return True78    except Exception as e:79        st.error(f"Erro ao salvar feedback: {str(e)}")80        return False81 82def fine_tune_model():83    """Realiza fine-tuning do modelo com dados de feedback positivo"""84    storage_path = get_storage_path()85    feedback_file = os.path.join(storage_path, 'feedback.json')86    87    if not os.path.exists(feedback_file):88        return89    90    try:91        # Carrega dados de feedback92        positive_examples = []93        with open(feedback_file, 'r') as f:94            for line in f:95                if line.strip():  # Ignora linhas vazias96                    feedback = json.loads(line)97                    if feedback['rating'] >= 4:  # Usa apenas feedback positivo98                        positive_examples.append({99                            'input_text': feedback['input_text'],100                            'output_text': feedback['output_text']101                        })102        103        if not positive_examples:104            return105        106        # Cria dataset e dataloader107        dataset = TextHumanizerDataset(positive_examples, st.session_state.tokenizer)108        dataloader = DataLoader(dataset, batch_size=4, shuffle=True)109        110        # Configura otimizador111        optimizer = torch.optim.AdamW(st.session_state.model.parameters(), lr=1e-5)112        113        # Fine-tuning114        st.session_state.model.train()115        for batch in dataloader:116            optimizer.zero_grad()117            outputs = st.session_state.model(118                input_ids=batch['input_ids'],119                attention_mask=batch['attention_mask'],120                labels=batch['labels']121            )122            loss = outputs.loss123            loss.backward()124            optimizer.step()125        126        st.session_state.model.eval()127        return True128    129    except Exception as e:130        st.error(f"Erro durante o fine-tuning: {str(e)}")131        return False132 133def clean_generated_text(text):134    """Remove comandos e limpa o texto gerado"""135    text = text.strip()136 137    # Lista de prefixos de comando para remover138    prefixes = [139        "reescreva o seguinte texto",140        "reescreva este texto",141        "reescreva o texto",142        "traduza",143        "humanize:",144        "humanizar:",145        "em português",146        "de forma mais natural"147    ]148 149    # Remove os prefixos de comando150    text_lower = text.lower()151    for prefix in prefixes:152        if text_lower.startswith(prefix):153            text = text[len(prefix):].strip()154            text_lower = text.lower()155    156    # Capitaliza a primeira letra157    if text:158        text = text[0].upper() + text[1:]159    160    return text161 162def humanize_text(text):163    """Humaniza o texto mantendo coerência e tamanho"""164    prompt = f"reescreva em português natural, mantendo todas as informações: {text}"165    166    # Tokenização com padding167    inputs = st.session_state.tokenizer(168        prompt,169        return_tensors="pt",170        max_length=512,  # Reduzido para evitar problemas de memória171        padding=True,172        truncation=True173    )174 175    # Parâmetros mais conservadores para geração176    try:177        outputs = st.session_state.model.generate(178            inputs.input_ids,179            max_length=512,  # Reduzido para maior estabilidade180            min_length=int(len(text.split()) * 0.8),  # Garante pelo menos 80% do tamanho original181            do_sample=False,  # Desativa amostragem para maior estabilidade182            num_beams=2,      # Reduzido para evitar problemas de memória183            repetition_penalty=1.1,  # Reduzido para evitar instabilidades184            length_penalty=1.0,      # Valor neutro185            early_stopping=True,     # Ativa early stopping186            no_repeat_ngram_size=2   # Evita repetições de bigramas187        )188        189        result = st.session_state.tokenizer.decode(outputs[0], skip_special_tokens=True)190        result = clean_generated_text(result)191        192        # Garante tamanho mínimo de forma mais suave193        if len(result.split()) < len(text.split()):194            missing_words = len(text.split()) - len(result.split())195            original_words = text.split()[-missing_words:]196            result = result + " " + " ".join(original_words)197        198        return result199        200    except Exception as e:201        st.error(f"Erro durante a geração: {str(e)}")202        # Fallback: retorna o texto original em caso de erro203        return text204 205# Initialize session state206if 'model_loaded' not in st.session_state:207    st.session_state.tokenizer = T5Tokenizer.from_pretrained("t5-base")208    st.session_state.model = T5ForConditionalGeneration.from_pretrained("t5-base")209    st.session_state.model_loaded = True210 211# UI Components212st.set_page_config(page_title="Advanced Text Humanizer", page_icon="🤖")213 214st.title("🤖 → 🧑 Humanizador de Texto Avançado")215st.markdown("""216Este aplicativo transforma textos robotizados em linguagem mais natural e humana, 217mantendo todas as informações originais e incluindo sistema de feedback para melhoria contínua.218""")219 220# Input area221input_text = st.text_area(222    "Cole seu texto de robô aqui:",223    height=150,224    help="Cole seu texto aqui para transformá-lo em uma versão mais natural e humana."225)226 227# Process button and results228if st.button("Humanizar", type="primary"):229    if not input_text:230        st.warning("⚠️ Por favor, cole um texto primeiro!")231    else:232        with st.spinner("Processando o texto..."):233            try:234                final_text = humanize_text(input_text)235                236                # Display results237                st.success("✨ Texto humanizado:")238                col1, col2 = st.columns(2)239                240                with col1:241                    st.text("Original:")242                    st.info(input_text)243                    st.write(f"Palavras: {len(input_text.split())}")244                245                with col2:246                    st.text("Resultado:")247                    st.info(final_text)248                    st.write(f"Palavras: {len(final_text.split())}")249                250                # Feedback section251                st.markdown("### Feedback")252                rating = st.slider(253                    "Como você avalia a qualidade do texto humanizado?",254                    min_value=1,255                    max_value=5,256                    value=3,257                    help="1 = Muito ruim, 5 = Excelente"258                )259                260                if st.button("Enviar Feedback"):261                    if save_feedback(input_text, final_text, rating):262                        st.success("Feedback salvo com sucesso! Obrigado pela contribuição.")263                        264                        # Trigger fine-tuning if we have enough positive feedback265                        if rating >= 4:266                            with st.spinner("Atualizando modelo com seu feedback..."):267                                if fine_tune_model():268                                    st.success("Modelo atualizado com sucesso!")269                                else:270                                    st.warning("Não foi possível atualizar o modelo neste momento.")271                    else:272                        st.error("Não foi possível salvar o feedback. Tente novamente mais tarde.")273                274            except Exception as e:275                st.error(f"❌ Erro no processamento: {str(e)}")276 277# Footer278st.markdown("---")279st.markdown(280    """281    <div style='text-align: center'>282        <small>Desenvolvido com ❤️ usando Streamlit e Transformers</small>283    </div>284    """,285    unsafe_allow_html=True286)