Update app.py

app.py

@@ -10,56 +10,51 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class DocumentGenerator:
-    """Gerencia a geração de documentos usando LLM"""
+    """Gerencia a geração de documentos usando modelo público"""
     
     def __init__(self):
-        # Usa a API key do ambiente
-        self.api = HfApi(token=os.environ.get("HF_TOKEN"))
-        # Usa modelo otimizado para português
-        self.model = "NeuralBeagle14-7B"  
+        # Usando um Space público que está efetivamente disponível
+        self.client = InferenceApi(
+            repo_id="OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+            token=os.environ.get("HF_TOKEN")
+        )
     
     def generate(self, doc_type: str, context: Dict[str, str]) -> str:
-        """Gera o documento usando LLM"""
+        """Gera o documento usando o modelo"""
         try:
-            base_prompt = f"""Atue como um advogado criminalista brasileiro experiente. 
-            Você deve gerar uma peça processual do tipo {doc_type} em português, formato jurídico brasileiro.
-            
-            DADOS:
-            - Cliente: {context.get('client_name')}
-            - Processo: {context.get('process_number')}
-            - Tribunal: {context.get('court')}
-            - Comarca: {context.get('jurisdiction')}
-            
-            FATOS:
-            {context.get('facts')}
-            
-            FUNDAMENTOS JURÍDICOS:
-            {context.get('legal_basis')}
-            
-            INSTRUÇÕES ESPECÍFICAS:
-            1. Use linguagem jurídica formal
-            2. Siga estritamente o formato de peças processuais brasileiras
-            3. Mantenha a estrutura com endereçamento, qualificação, fatos, direito e pedidos
-            4. Inclua local e data ao final
-            5. Use apenas os dados fornecidos
-            6. Evite repetições desnecessárias
-            7. Cite corretamente artigos e jurisprudência
-            8. Use formatação profissional
-            
-            COMECE A GERAR O DOCUMENTO AGORA:"""
+            prompt = f"""You are a Brazilian criminal lawyer. Create a {doc_type} in Portuguese following Brazilian legal standards.
+
+Information:
+Client: {context.get('client_name')}
+Process: {context.get('process_number')}
+Court: {context.get('court')}
+Jurisdiction: {context.get('jurisdiction')}
 
-            # Faz a chamada à API
-            inference = InferenceApi(repo_id=f"meta-llama/{self.model}", token=os.environ.get("HF_TOKEN"))
-            response = inference(
-                base_prompt,
-                max_length=4096,
-                temperature=0.3,
-                top_p=0.95,
-                repetition_penalty=1.15,
-                do_sample=True
+Facts:
+{context.get('facts')}
+
+Legal Basis:
+{context.get('legal_basis')}
+
+Instructions:
+1. Use formal legal Portuguese
+2. Follow Brazilian legal document format
+3. Include all required sections
+4. Be precise and clear
+5. Keep proper legal formatting"""
+
+            response = self.client(
+                inputs=prompt,
+                parameters={
+                    "max_new_tokens": 2048,
+                    "temperature": 0.3,
+                    "top_p": 0.95,
+                    "repetition_penalty": 1.15,
+                    "do_sample": True
+                }
             )
 
-            return self._format_output(response[0]['generated_text'])
+            return self._format_output(response[0]["generated_text"])
             
         except Exception as e:
             logger.error(f"Erro na geração: {str(e)}")
@@ -69,18 +64,17 @@ class DocumentGenerator:
         """Formata o texto gerado"""
         if not text:
             return "Erro: Nenhum texto gerado"
-
+        
         # Remove o prompt da resposta
-        text = text.split("COMECE A GERAR O DOCUMENTO AGORA:")[-1].strip()
+        text = text.split("Instructions:")[-1].strip()
         
         # Ajusta formatação
         lines = [line.strip() for line in text.split('\n') if line.strip()]
         formatted_text = '\n\n'.join(lines)
         
-        # Adiciona data se não existir
-        if "[DATA]" in formatted_text:
-            current_date = datetime.now().strftime('%d de %B de %Y')
-            formatted_text = formatted_text.replace("[DATA]", current_date)
+        # Adiciona data atual
+        current_date = datetime.now().strftime('%d de %B de %Y')
+        formatted_text = f"{formatted_text}\n\n{context.get('jurisdiction')}, {current_date}"
         
         return formatted_text