Update app.py

app.py

@@ -5,6 +5,7 @@ from torch.utils.data import Dataset, DataLoader
 import json
 import os
 from datetime import datetime
+import tempfile
 
 # Custom dataset for fine-tuning
 class TextHumanizerDataset(Dataset):
@@ -40,6 +41,17 @@ class TextHumanizerDataset(Dataset):
             'labels': target_encoding['input_ids'].squeeze()
         }
 
+def get_storage_path():
+    """Retorna o caminho correto para armazenamento no Hugging Face Spaces"""
+    if os.environ.get('SPACE_ID'):  # Verifica se está rodando no Spaces
+        return '/data'  # Diretório persistente no Spaces
+    else:
+        # Fallback para desenvolvimento local
+        temp_dir = tempfile.gettempdir()
+        feedback_dir = os.path.join(temp_dir, 'feedback_data')
+        os.makedirs(feedback_dir, exist_ok=True)
+        return feedback_dir
+
 def save_feedback(input_text, output_text, rating):
     """Salva o feedback do usuário para futuro treinamento"""
     feedback_data = {
@@ -49,53 +61,74 @@ def save_feedback(input_text, output_text, rating):
         'timestamp': datetime.now().isoformat()
     }
     
-    # Cria diretório se não existir
-    os.makedirs('feedback_data', exist_ok=True)
+    storage_path = get_storage_path()
+    feedback_file = os.path.join(storage_path, 'feedback.json')
     
-    # Salva em arquivo JSON
-    with open('feedback_data/feedback.json', 'a') as f:
-        f.write(json.dumps(feedback_data) + '\n')
+    try:
+        # Cria arquivo se não existir
+        if not os.path.exists(feedback_file):
+            with open(feedback_file, 'w') as f:
+                f.write('')
+        
+        # Append do novo feedback
+        with open(feedback_file, 'a') as f:
+            f.write(json.dumps(feedback_data) + '\n')
+            
+        return True
+    except Exception as e:
+        st.error(f"Erro ao salvar feedback: {str(e)}")
+        return False
 
 def fine_tune_model():
     """Realiza fine-tuning do modelo com dados de feedback positivo"""
-    if not os.path.exists('feedback_data/feedback.json'):
-        return
-    
-    # Carrega dados de feedback
-    positive_examples = []
-    with open('feedback_data/feedback.json', 'r') as f:
-        for line in f:
-            feedback = json.loads(line)
-            if feedback['rating'] >= 4:  # Usa apenas feedback positivo
-                positive_examples.append({
-                    'input_text': feedback['input_text'],
-                    'output_text': feedback['output_text']
-                })
+    storage_path = get_storage_path()
+    feedback_file = os.path.join(storage_path, 'feedback.json')
     
-    if not positive_examples:
+    if not os.path.exists(feedback_file):
         return
     
-    # Cria dataset e dataloader
-    dataset = TextHumanizerDataset(positive_examples, st.session_state.tokenizer)
-    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
-    
-    # Configura otimizador
-    optimizer = torch.optim.AdamW(st.session_state.model.parameters(), lr=1e-5)
-    
-    # Fine-tuning
-    st.session_state.model.train()
-    for batch in dataloader:
-        optimizer.zero_grad()
-        outputs = st.session_state.model(
-            input_ids=batch['input_ids'],
-            attention_mask=batch['attention_mask'],
-            labels=batch['labels']
-        )
-        loss = outputs.loss
-        loss.backward()
-        optimizer.step()
+    try:
+        # Carrega dados de feedback
+        positive_examples = []
+        with open(feedback_file, 'r') as f:
+            for line in f:
+                if line.strip():  # Ignora linhas vazias
+                    feedback = json.loads(line)
+                    if feedback['rating'] >= 4:  # Usa apenas feedback positivo
+                        positive_examples.append({
+                            'input_text': feedback['input_text'],
+                            'output_text': feedback['output_text']
+                        })
+        
+        if not positive_examples:
+            return
+        
+        # Cria dataset e dataloader
+        dataset = TextHumanizerDataset(positive_examples, st.session_state.tokenizer)
+        dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
+        
+        # Configura otimizador
+        optimizer = torch.optim.AdamW(st.session_state.model.parameters(), lr=1e-5)
+        
+        # Fine-tuning
+        st.session_state.model.train()
+        for batch in dataloader:
+            optimizer.zero_grad()
+            outputs = st.session_state.model(
+                input_ids=batch['input_ids'],
+                attention_mask=batch['attention_mask'],
+                labels=batch['labels']
+            )
+            loss = outputs.loss
+            loss.backward()
+            optimizer.step()
+        
+        st.session_state.model.eval()
+        return True
     
-    st.session_state.model.eval()
+    except Exception as e:
+        st.error(f"Erro durante o fine-tuning: {str(e)}")
+        return False
 
 def clean_generated_text(text):
     """Remove comandos e limpa o texto gerado"""
@@ -130,33 +163,44 @@ def humanize_text(text):
     """Humaniza o texto mantendo coerência e tamanho"""
     prompt = f"reescreva em português natural, mantendo todas as informações: {text}"
     
-    input_ids = st.session_state.tokenizer(
+    # Tokenização com padding
+    inputs = st.session_state.tokenizer(
         prompt,
         return_tensors="pt",
-        max_length=512,
+        max_length=512,  # Reduzido para evitar problemas de memória
+        padding=True,
         truncation=True
-    ).input_ids
-
-    # Parâmetros ajustados para melhor coerência
-    outputs = st.session_state.model.generate(
-        input_ids,
-        max_length=1024,
-        min_length=len(text.split()),
-        do_sample=False,
-        temperature=0.1,
-        top_p=0.95,
-        num_beams=2,
-        repetition_penalty=1.1,
-        length_penalty=1.0
     )
-    result = st.session_state.tokenizer.decode(outputs[0], skip_special_tokens=True)
-    result = clean_generated_text(result)
-    
-    # Garante tamanho mínimo
-    while len(result.split()) < len(text.split()):
-        result += " " + " ".join(text.split()[-(len(text.split()) - len(result.split())):])
-    
-    return result
+
+    # Parâmetros mais conservadores para geração
+    try:
+        outputs = st.session_state.model.generate(
+            inputs.input_ids,
+            max_length=512,  # Reduzido para maior estabilidade
+            min_length=int(len(text.split()) * 0.8),  # Garante pelo menos 80% do tamanho original
+            do_sample=False,  # Desativa amostragem para maior estabilidade
+            num_beams=2,      # Reduzido para evitar problemas de memória
+            repetition_penalty=1.1,  # Reduzido para evitar instabilidades
+            length_penalty=1.0,      # Valor neutro
+            early_stopping=True,     # Ativa early stopping
+            no_repeat_ngram_size=2   # Evita repetições de bigramas
+        )
+        
+        result = st.session_state.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        result = clean_generated_text(result)
+        
+        # Garante tamanho mínimo de forma mais suave
+        if len(result.split()) < len(text.split()):
+            missing_words = len(text.split()) - len(result.split())
+            original_words = text.split()[-missing_words:]
+            result = result + " " + " ".join(original_words)
+        
+        return result
+        
+    except Exception as e:
+        st.error(f"Erro durante a geração: {str(e)}")
+        # Fallback: retorna o texto original em caso de erro
+        return text
 
 # Initialize session state
 if 'model_loaded' not in st.session_state:
@@ -214,14 +258,18 @@ if st.button("Humanizar", type="primary"):
                 )
                 
                 if st.button("Enviar Feedback"):
-                    save_feedback(input_text, final_text, rating)
-                    st.success("Feedback salvo com sucesso! Obrigado pela contribuição.")
-                    
-                    # Trigger fine-tuning if we have enough positive feedback
-                    if rating >= 4:
-                        with st.spinner("Atualizando modelo com seu feedback..."):
-                            fine_tune_model()
-                        st.success("Modelo atualizado com sucesso!")
+                    if save_feedback(input_text, final_text, rating):
+                        st.success("Feedback salvo com sucesso! Obrigado pela contribuição.")
+                        
+                        # Trigger fine-tuning if we have enough positive feedback
+                        if rating >= 4:
+                            with st.spinner("Atualizando modelo com seu feedback..."):
+                                if fine_tune_model():
+                                    st.success("Modelo atualizado com sucesso!")
+                                else:
+                                    st.warning("Não foi possível atualizar o modelo neste momento.")
+                    else:
+                        st.error("Não foi possível salvar o feedback. Tente novamente mais tarde.")
                 
             except Exception as e:
                 st.error(f"❌ Erro no processamento: {str(e)}")