Update analyzers/ner_analyzer.py

analyzers/ner_analyzer.py

@@ -2,7 +2,7 @@ from transformers import AutoModelForTokenClassification, AutoTokenizer
 import torch
 from typing import List, Tuple
 import logging
-import re
+import spacy
 from .base_analyzer import BaseAnalyzer
 
 logger = logging.getLogger(__name__)
@@ -14,6 +14,9 @@ class NERAnalyzer(BaseAnalyzer):
         self.model = AutoModelForTokenClassification.from_pretrained(self.model_name)
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         logger.info("Modelo NER e tokenizador carregados com sucesso")
+        
+        # Carregar modelo spaCy para processamento de dependências e identificação de entidades
+        self.nlp = spacy.load("pt_core_news_lg")
 
     def extract_entities(self, text: str) -> List[Tuple[str, str]]:
         logger.debug("Iniciando extração de entidades com NER")
@@ -53,20 +56,36 @@ class NERAnalyzer(BaseAnalyzer):
         
         return representatives
 
-    def extract_participation_percentage(self, text: str) -> dict:
-        # Usando expressão regular para procurar percentuais de participação
+    def extract_participation_percentage(self, text: str, representatives: List[str]) -> dict:
+        """
+        Identifica os percentuais de participação no texto e associa ao representante mais próximo.
+        """
+        # Utilizando o spaCy para análise de dependências e encontrar a relação entre porcentagens e nomes
+        doc = self.nlp(text)
         participation_data = {}
-        
-        # Expressão regular para capturar percentuais no formato numérico seguido de "%"
-        pattern = re.compile(r"(\d{1,3},\d{1,2})%\s*(de\s+)?([A-Za-z\s]+)")
-        
-        # Procurando no texto os percentuais e associando-os aos nomes
-        for match in pattern.finditer(text):
-            percentage = float(match.group(1).replace(',', '.'))  # Convertendo percentual para float
-            name = match.group(3).strip()  # Nome do sócio ou pessoa
-            
-            participation_data[name] = percentage
-        
+
+        # Iterar sobre o texto para buscar percentuais e associá-los ao nome mais próximo
+        for token in doc:
+            # Verificar se o token é um percentual (%)
+            if token.text.endswith('%'):
+                percentage = token.text.replace('%', '').strip()
+                try:
+                    percentage = float(percentage.replace(',', '.'))  # Garantir que o valor seja numérico
+                except ValueError:
+                    continue
+
+                # Verificar os nomes próximos ao percentual
+                left_context = [w.text for w in token.lefts]
+                right_context = [w.text for w in token.rights]
+                context = left_context + right_context
+
+                # Procurar por um representante no contexto
+                for rep in representatives:
+                    # Se o nome do representante estiver nas palavras ao redor do percentual, associamos
+                    if any(rep.lower() in word.lower() for word in context):
+                        participation_data[rep] = percentage
+                        break
+
         logger.debug(f"Dados de participação: {participation_data}")
         return participation_data
 
@@ -75,23 +94,12 @@ class NERAnalyzer(BaseAnalyzer):
         entities = self.extract_entities(text)
         # Extraímos os representantes a partir das entidades
         representatives = self.extract_representatives(entities)
-        # Extraímos os percentuais de participação
-        participation_data = self.extract_participation_percentage(text)
+        # Extraímos os percentuais de participação e associamos aos representantes
+        participation_data = self.extract_participation_percentage(text, representatives)
         
         # Associa os representantes aos seus percentuais de participação, se disponível
         representatives_with_percentage = []
         for rep in representatives:
             # Tentando associar o nome do representante ao percentual de participação
             if rep in participation_data:
-                representatives_with_percentage.append(f"{rep} ({participation_data[rep]}%)")
-            else:
-                representatives_with_percentage.append(rep)
-        
-        return representatives_with_percentage
-
-    def format_output(self, representatives: List[str]) -> str:
-        output = "ANÁLISE DO CONTRATO SOCIAL (NER)\n\n"
-        output += "REPRESENTANTES IDENTIFICADOS:\n"
-        for rep in representatives:
-            output += f"- {rep}\n"
-        return output
+