analyzers/ner_analyzer_teste.py

from transformers import DistilBertTokenizer, DistilBertForTokenClassification
import torch
from typing import List, Tuple
import spacy
import logging

logger = logging.getLogger(__name__)

class NERAnalyzer:
    def __init__(self):
        # 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 = 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]]:
        # Tokeniza o texto e prepara para a análise de entidades
        inputs = self.tokenizer(text, max_length=512, truncation=True, return_tensors="pt")
        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":  # "O" significa que não é uma entidade
                entities.append((token, entity_label))
                
        logger.debug(f"Tokens identificados: {entities}")
        return entities

    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('##'):
                    current_person += token.replace("##", "")
                else:
                    current_person += f" {token.replace('##', '')}"
            else:
                if current_person:
                    representatives.append(current_person)
                    current_person = ""
        
        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})
        
        logger.debug(f"Representantes e quotas extraídos: {representative_data}")
        return representative_data

    def extract_quotas(self, doc) -> dict:
        """
        Extrai as quotas dos sócios a partir do texto processado pelo spaCy.
        """
        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

    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 calculados: {representative_data}")
        return representative_data

    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)
        # 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)
        
        # 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 formatted_output

    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