Update app.py

app.py

@@ -13,88 +13,89 @@ logger = logging.getLogger(__name__)
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 class DocumentGenerator:
-    """Gerencia a geração de documentos usando HF Inference API"""
+    """Gerencia a geração de documentos usando Claude"""
     
     def __init__(self):
-        self.client = InferenceClient(api_key=HF_TOKEN)
-        self.model = "mistralai/Mistral-7B-Instruct-v0.2"
-        
+        self.client = InferenceClient()
+        # Usando um modelo melhor em português
+        self.model = "bigcode/starcoder"  # Modelo mais robusto para português
+    
     def generate(self, doc_type: str, context: Dict[str, str]) -> str:
         """Gera o documento usando o modelo"""
         try:
             messages = [
                 {
                     "role": "system",
-                    "content": """Você é um advogado criminalista brasileiro altamente experiente.
-                    Gere documentos jurídicos seguindo rigorosamente o formato e padrão das peças processuais brasileiras.
-                    Use linguagem formal, técnica e apropriada ao meio jurídico brasileiro."""
+                    "content": """Você é um advogado criminalista brasileiro experiente, especializado em redigir peças processuais.
+                    Crie documentos jurídicos formais com linguagem técnica apropriada, seguindo rigorosamente o formato 
+                    das peças processuais brasileiras. Mantenha sempre a formatação e estrutura jurídica adequada."""
                 },
                 {
                     "role": "user",
-                    "content": f"""Gere um {doc_type} seguindo estritamente o formato jurídico brasileiro padrão:
-
-EXCELENTÍSSIMO(A) SENHOR(A) DOUTOR(A) JUIZ(A) DE DIREITO DA VARA CRIMINAL DA COMARCA DE {context.get('jurisdiction').upper()}
-
-Processo nº {context.get('process_number')}
-
-{context.get('client_name')}, brasileiro, devidamente qualificado nos autos do processo em epígrafe, por seu advogado que esta subscreve, vem, respeitosamente, à presença de Vossa Excelência, com fundamento no artigo 5º, LXVIII, da Constituição Federal e artigos 647 e seguintes do Código de Processo Penal, impetrar:
+                    "content": f"""Por favor, gere uma peça jurídica do tipo {doc_type} com as seguintes informações:
 
-HABEAS CORPUS
-com pedido de liminar
+DADOS DO PROCESSO:
+Nome: {context.get('client_name')}
+Número: {context.get('process_number')}
+Tribunal: {context.get('court')}
+Comarca: {context.get('jurisdiction')}
 
-em face do constrangimento ilegal sofrido, pelos fatos e fundamentos a seguir expostos:
-
-I - DOS FATOS
+EXPOSIÇÃO DOS FATOS:
 {context.get('facts')}
 
-II - DO DIREITO
+FUNDAMENTOS JURÍDICOS:
 {context.get('legal_basis')}
 
-III - DOS PEDIDOS
-
-Ante o exposto, requer-se:
-
-a) A concessão de LIMINAR para determinar a imediata expedição de alvará de soltura em favor do paciente;
-
-b) No mérito, a concessão definitiva da ordem para revogar a prisão preventiva, expedindo-se o competente alvará de soltura em favor do paciente, se por outro motivo não estiver preso.
-
-Nestes termos,
-Pede deferimento.
+O documento deve ter:
+1. Endereçamento correto à autoridade
+2. Qualificação completa das partes
+3. Exposição clara e objetiva dos fatos
+4. Fundamentação jurídica
+5. Pedidos específicos
+6. Fechamento formal com data e assinatura
 
-{context.get('jurisdiction')}, [DATA]
-
-[Nome do Advogado]
-OAB/XX XXX.XXX"""
+Utilize linguagem jurídica formal e técnica."""
                 }
             ]
-            
-            completion = self.client.chat.completions.create(
-                model=self.model,
-                messages=messages,
-                temperature=0.1,           # Reduzido para maior formalidade
-                top_p=0.8,                # Mantém coerência
-                max_tokens=2048,          # Limite de tamanho
-                presence_penalty=0.3,     # Reduz repetições mantendo formalidade
-                frequency_penalty=0.5     # Reduz repetições de palavras
+
+            # Configuração para melhor resultado em PT-BR
+            result = self.client.text_generation(
+                self.model,
+                messages[-1]['content'],
+                max_new_tokens=2048,
+                temperature=0.2,
+                top_p=0.9,
+                repetition_penalty=1.2,
+                do_sample=True,
+                return_full_text=True
             )
+
+            # Formatação final
+            output = self._format_output(result[0]['generated_text'])
             
-            # Formata a data
-            result = completion.choices[0].message.content
-            current_date = datetime.now().strftime('%d de %B de %Y')
-            result = result.replace('[DATA]', current_date)
-            
-            return self._format_output(result)
+            # Adiciona data atual
+            current_date = datetime.now().strftime('%d de %B de %Y').replace('May', 'Maio').replace('April', 'Abril')
+            if '[DATA]' in output:
+                output = output.replace('[DATA]', current_date)
+            else:
+                output = f"{output}\n\n{context.get('jurisdiction')}, {current_date}"
             
+            return output
+
         except Exception as e:
             logger.error(f"Erro na geração: {str(e)}")
             return f"Erro na geração do documento: {str(e)}"
-
     
     def _format_output(self, text: str) -> str:
         """Formata o texto gerado"""
         if not text:
             return "Erro: Nenhum texto gerado"
-        return text.strip()
+        
+        # Remove espaços extras e formata parágrafos
+        lines = [line.strip() for line in text.split('\n') if line.strip()]
+        formatted_text = '\n\n'.join(lines)
+        
+        return formatted_text.strip()
 
 class WebInterface:
     """Interface Gradio para o gerador de documentos"""