Update app.py

app.py

@@ -93,36 +93,77 @@ class PromptManager:
         return prompts.get(doc_type, "Tipo de documento não suportado")
 
 class DocumentGenerator:
-    """Gerencia a geração de documentos usando GPT4All-J"""
+    """Gerencia a geração de documentos usando OpenHermes"""
     
     def __init__(self):
-        self.pipe = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
+        self.pipe = pipeline(
+            "text-generation",
+            model="jondurbin/airoboros-l2-70b-gpt4-1.4.1",  # Modelo público com bom desempenho em PT
+            trust_remote_code=True
+        )
         self.prompt_manager = PromptManager()
     
     def generate(self, doc_type: str, context: Dict[str, str]) -> str:
         """Gera o documento baseado no tipo e contexto"""
         try:
-            prompt = self.prompt_manager.get_base_prompt(doc_type, context)
+            # Prepara o prompt com instruções específicas
+            base_prompt = self.prompt_manager.get_base_prompt(doc_type, context)
             
-            # Configuração para a geração
+            system_prompt = """Você é um advogado criminalista experiente especializado em gerar documentos jurídicos.
+            Gere o documento solicitado usando linguagem formal e técnica apropriada."""
+            
+            formatted_prompt = f"""### System:
+            {system_prompt}
+            
+            ### User:
+            {base_prompt}
+            
+            ### Assistant:
+            Gerando o documento solicitado em formato jurídico adequado:
+            """
+            
+            # Configuração otimizada para documentos jurídicos
             result = self.pipe(
-                prompt,
-                max_length=2048,
-                temperature=0.7,
-                top_p=0.95,
+                formatted_prompt,
+                max_new_tokens=2048,          # Aumentado para documentos longos
+                temperature=0.3,              # Reduzido para maior consistência
+                top_p=0.85,                   # Ajustado para melhor coerência
+                top_k=40,                     # Limita as opções de tokens
                 num_return_sequences=1,
-                do_sample=True
+                do_sample=True,
+                repetition_penalty=1.2,       # Evita repetições
+                pad_token_id=self.pipe.tokenizer.eos_token_id
             )[0]['generated_text']
             
-            return self._format_output(result)
+            # Remove o prompt do resultado e retorna apenas o documento
+            final_text = result.split("### Assistant:")[-1].strip()
+            return self._format_output(final_text)
             
         except Exception as e:
             logger.error(f"Erro na geração: {str(e)}")
-            return f"Erro na geração do documento: {str(e)}"
+            if "CUDA" in str(e):
+                return "Erro: Problema com recursos GPU. Tentando novamente com configurações reduzidas..."
+            elif "out of memory" in str(e):
+                return "Erro: Memória insuficiente. Tente um documento menor."
+            else:
+                return f"Erro na geração do documento: {str(e)}"
     
     def _format_output(self, text: str) -> str:
         """Formata o texto gerado para melhor apresentação"""
-        return text.strip()
+        if not isinstance(text, str):
+            if isinstance(text, (list, tuple)) and len(text) > 0:
+                text = str(text[0])
+            else:
+                text = str(text)
+                
+        # Remove linhas vazias extras
+        lines = [line for line in text.split('\n') if line.strip()]
+        formatted_text = '\n'.join(lines)
+        
+        # Ajusta espaçamento de parágrafos
+        formatted_text = formatted_text.replace('\n\n\n', '\n\n')
+        
+        return formatted_text.strip()
 
 class WebInterface:
     """Interface Gradio para o gerador de documentos"""