Update app.py

app.py

@@ -1,27 +1,29 @@
 import gradio as gr
-from transformers import pipeline
-import json
+from huggingface_hub import InferenceClient
 from datetime import datetime
 import logging
 from typing import Dict, List, Optional
 
 # Configuração do logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-class PromptManager:
-    """Gerencia a criação e otimização de prompts"""
+class DocumentGenerator:
+    """Gerencia a geração de documentos usando Mistral via InferenceClient"""
     
-    @staticmethod
-    def get_base_prompt(doc_type: str, context: Dict[str, str]) -> str:
-        """Retorna o prompt base para cada tipo de documento"""
-        
-        prompts = {
-            "habeas_corpus": f"""
-            Atue como um advogado criminalista experiente e gere um Habeas Corpus detalhado:
+    def __init__(self):
+        # Inicializa o cliente de inferência
+        self.client = InferenceClient()
+        self.model = "mistralai/Mistral-7B-Instruct-v0.2"
+    
+    def generate(self, doc_type: str, context: Dict[str, str]) -> str:
+        """Gera o documento baseado no tipo e contexto"""
+        try:
+            # Prepara a mensagem para o modelo
+            system_prompt = """Você é um advogado criminalista experiente especializado em gerar documentos jurídicos no Brasil.
+            Gere o documento solicitado usando linguagem formal, técnica e no formato jurídico correto."""
+            
+            user_content = f"""Gere um {doc_type} com os seguintes dados:
             
             DADOS DO CASO:
             Cliente: {context.get('client_name')}
@@ -35,135 +37,33 @@ class PromptManager:
             FUNDAMENTOS JURÍDICOS:
             {context.get('legal_basis')}
             
-            Gere um Habeas Corpus completo e tecnicamente preciso, incluindo:
-            1. Cabeçalho com endereçamento correto
-            2. Qualificação das partes
-            3. Fatos relevantes
-            4. Fundamentos jurídicos
-            5. Jurisprudência relevante
-            6. Pedidos (liminar e mérito)
-            """,
-            
-            "denuncia": f"""
-            Atue como um promotor de justiça experiente e gere uma denúncia criminal:
-            
-            DADOS DO CASO:
-            Processo: {context.get('process_number')}
-            Tribunal: {context.get('court')}
-            Comarca: {context.get('jurisdiction')}
-            
-            FATOS:
-            {context.get('facts')}
-            
-            ENQUADRAMENTO LEGAL:
-            {context.get('legal_basis')}
-            
-            Gere uma denúncia criminal completa, incluindo:
-            1. Qualificação do(s) denunciado(s)
-            2. Narrativa clara dos fatos
-            3. Tipificação penal
-            4. Circunstâncias do crime
-            5. Rol de testemunhas
-            6. Pedidos
-            """,
-            
-            "alegacoes_finais": f"""
-            Atue como um advogado criminalista e gere alegações finais:
-            
-            DADOS DO PROCESSO:
-            Cliente: {context.get('client_name')}
-            Processo: {context.get('process_number')}
-            Tribunal: {context.get('court')}
-            
-            FATOS E PROVAS:
-            {context.get('facts')}
-            
-            TESE DEFENSIVA:
-            {context.get('legal_basis')}
-            
-            Elabore alegações finais completas, incluindo:
-            1. Relatório do processo
-            2. Preliminares (se houver)
-            3. Mérito
-            4. Análise das provas
-            5. Pedidos
-            """
-        }
-        
-        return prompts.get(doc_type, "Tipo de documento não suportado")
+            O documento deve ser completo, tecnicamente preciso e seguir o formato jurídico brasileiro."""
 
-class DocumentGenerator:
-    """Gerencia a geração de documentos usando OpenHermes"""
-    
-    def __init__(self):
-        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:
-            # Prepara o prompt com instruções específicas
-            base_prompt = self.prompt_manager.get_base_prompt(doc_type, context)
-            
-            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(
-                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,
-                repetition_penalty=1.2,       # Evita repetições
-                pad_token_id=self.pipe.tokenizer.eos_token_id
-            )[0]['generated_text']
-            
-            # Remove o prompt do resultado e retorna apenas o documento
-            final_text = result.split("### Assistant:")[-1].strip()
-            return self._format_output(final_text)
+            messages = [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_content}
+            ]
+
+            # Faz a chamada ao modelo
+            completion = self.client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                max_tokens=2048,
+                temperature=0.3,
+                top_p=0.85,
+            )
+
+            return self._format_output(completion.choices[0].message.content)
             
         except Exception as e:
             logger.error(f"Erro na geração: {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)}"
+            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"""
-        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()
+        """Formata o texto gerado"""
+        if not text:
+            return "Erro: Nenhum texto gerado"
+        return text.strip()
 
 class WebInterface:
     """Interface Gradio para o gerador de documentos"""
@@ -178,7 +78,7 @@ class WebInterface:
         with gr.Blocks(theme=gr.themes.Soft()) as self.app:
             gr.Markdown("""
             # Gerador de Documentos Criminais
-            ### Powered by GPT4All-J
+            ### Powered by Mistral-7B-Instruct
             """)
             
             with gr.Row():
@@ -188,7 +88,10 @@ class WebInterface:
                         choices=[
                             "Habeas Corpus",
                             "Denúncia",
-                            "Alegações Finais"
+                            "Alegações Finais",
+                            "Resposta à Acusação",
+                            "Recurso em Sentido Estrito",
+                            "Apelação Criminal"
                         ],
                         label="Tipo de Documento",
                         value="Habeas Corpus"
@@ -228,12 +131,6 @@ class WebInterface:
                             placeholder="Indique os fundamentos legais principais..."
                         )
                     
-                    # Status e botões
-                    status = gr.Textbox(
-                        label="Status",
-                        interactive=False
-                    )
-                    
                     generate_btn = gr.Button(
                         "Gerar Documento",
                         variant="primary"
@@ -245,8 +142,13 @@ class WebInterface:
                         lines=30,
                         show_copy_button=True
                     )
+                    
+                    status = gr.Textbox(
+                        label="Status",
+                        interactive=False
+                    )
             
-            # Exemplos de uso
+            # Exemplo de uso
             gr.Examples(
                 examples=[
                     [
@@ -302,7 +204,7 @@ class WebInterface:
             
             # Gera documento
             result = self.generator.generate(
-                doc_type.lower().replace(" ", "_"),
+                doc_type,
                 context
             )
             
@@ -314,11 +216,8 @@ class WebInterface:
     
     def launch(self):
         """Inicia a interface web"""
-        self.app.launch(share=True)
+        self.app.launch()
 
 if __name__ == "__main__":
-    # Dependências necessárias:
-    # pip install gradio transformers torch
-    
     interface = WebInterface()
     interface.launch()