Update analyzers/ner_analyzer.py

analyzers/ner_analyzer.py

@@ -1,45 +1,50 @@
-from transformers import AutoModelForTokenClassification, AutoTokenizer
+from transformers import DistilBertTokenizer, DistilBertForTokenClassification
 import torch
 from typing import List, Tuple
-import logging
 import spacy
-from .base_analyzer import BaseAnalyzer
+import logging
 
 logger = logging.getLogger(__name__)
 
-class NERAnalyzer(BaseAnalyzer):
+class NERAnalyzer:
     def __init__(self):
-        self.model_name = "dominguesm/ner-legal-bert-base-cased-ptbr"
+        # Usando DistilBERT para token classification
+        self.model_name = "dbmdz/bert-large-cased-finetuned-conll03-english"
         logger.info(f"Carregando o modelo NER: {self.model_name}")
-        self.model = AutoModelForTokenClassification.from_pretrained(self.model_name)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = DistilBertForTokenClassification.from_pretrained(self.model_name)
+        self.tokenizer = DistilBertTokenizer.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")
+        # Tokeniza o texto e prepara para a análise de entidades
         inputs = self.tokenizer(text, max_length=512, truncation=True, return_tensors="pt")
-        tokens = inputs.tokens()
+        tokens = self.tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
         
         with torch.no_grad():
             outputs = self.model(**inputs).logits
             predictions = torch.argmax(outputs, dim=2)
         
+        # Mapeia as predições para labels de entidades
         entities = []
         for token, prediction in zip(tokens, predictions[0].numpy()):
             entity_label = self.model.config.id2label[prediction]
-            if entity_label != "O":
+            if entity_label != "O":  # "O" significa que não é uma entidade
                 entities.append((token, entity_label))
                 
-        logger.info(f"tokens: {entities}")
+        logger.debug(f"Tokens identificados: {entities}")
         return entities
 
-    def extract_representatives(self, entities: List[Tuple[str, str]]) -> List[str]:
+    def extract_representatives_and_quotas(self, entities: List[Tuple[str, str]], text: str) -> List[dict]:
+        """
+        Extrai os representantes legais e as quotas a partir do texto usando o modelo NER.
+        """
         representatives = []
         current_person = ""
         
+        # Encontrar os sócios a partir das entidades
         for token, label in entities:
             if label in ["B-PESSOA", "I-PESSOA"]:
                 if token.startswith('##'):
@@ -53,60 +58,72 @@ class NERAnalyzer(BaseAnalyzer):
         
         if current_person:
             representatives.append(current_person)
+
+        # Agora, vamos analisar o texto para encontrar as quotas associadas aos sócios
+        doc = self.nlp(text)
+        quota_values = self.extract_quotas(doc)
+        
+        # Associa os representantes com suas quotas
+        representative_data = []
+        for rep in representatives:
+            if rep in quota_values:
+                representative_data.append({"representante": rep, "quotas": quota_values[rep]})
+            else:
+                representative_data.append({"representante": rep, "quotas": 0})
         
-        return representatives
+        logger.debug(f"Representantes e quotas extraídos: {representative_data}")
+        return representative_data
 
-    def extract_participation_percentage(self, text: str, representatives: List[str]) -> dict:
+    def extract_quotas(self, doc) -> dict:
         """
-        Identifica os percentuais de participação no texto e associa ao representante mais próximo.
+        Extrai as quotas dos sócios a partir do texto processado pelo spaCy.
         """
-        # Utilizando o spaCy para análise de dependências e encontrar a relação entre porcentagens e nomes
-        doc = self.nlp(text)
-        participation_data = {}
-
-        # 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
+        quota_values = {}
+        # Buscando por padrões relacionados a quotas utilizando as dependências sintáticas
+        for ent in doc.ents:
+            if ent.label_ == "MONEY" and 'quota' in ent.sent.text.lower():
+                # Encontrar o sócio associado à quota
+                for token in ent.sent:
+                    if token.dep_ == "nsubj" and token.pos_ == "PROPN":
+                        # A entidade que está associada à quota é o sujeito da frase
+                        name = token.text
+                        if name not in quota_values:
+                            quota_values[name] = 0
+                        # Adicionar a quota à pessoa associada
+                        quota_values[name] += float(ent.text.replace("R$", "").replace(",", ".").strip())
+        
+        logger.debug(f"Valores de quotas extraídos com o spaCy: {quota_values}")
+        return quota_values
 
-                # 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
+    def calculate_participation(self, total_quotas: int, total_capital: float, representative_data: List[dict]) -> List[dict]:
+        """
+        Calcula a participação de cada sócio com base nas quotas e no capital total.
+        """
+        quota_value = total_capital / total_quotas  # Valor de cada quota
+        for data in representative_data:
+            quotas = data["quotas"]
+            percentage = (quotas / total_quotas) * 100
+            data["percentual"] = round(percentage, 2)
+            data["valor"] = quotas * quota_value
 
-        logger.debug(f"Dados de participação: {participation_data}")
-        return participation_data
+        logger.debug(f"Dados de participação calculados: {representative_data}")
+        return representative_data
 
-    def analyze(self, text: str) -> List[str]:
-        # Extraímos as entidades (nomes) do texto
+    def analyze(self, text: str, total_quotas: int, total_capital: float) -> List[str]:
+        # Passo 1: Extrair as entidades (nomes dos sócios) do texto
         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 e associamos aos representantes
-        participation_data = self.extract_participation_percentage(text, representatives)
+        # Passo 2: Extrair representantes e associá-los com quotas
+        representative_data = self.extract_representatives_and_quotas(entities, text)
+        # Passo 3: Calcular a participação de cada representante com base nas quotas
+        participation_data = self.calculate_participation(total_quotas, total_capital, representative_data)
         
-        # 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)
+        # Formatar a saída final com representantes e seus percentuais
+        formatted_output = []
+        for data in participation_data:
+            rep = data["representante"]
+            formatted_output.append(f"{rep} - {data['percentual']}% (R${data['valor']})")
         
-        return representatives_with_percentage
+        return formatted_output
 
     def format_output(self, representatives: List[str]) -> str:
         output = "ANÁLISE DO CONTRATO SOCIAL (NER)\n\n"