model_crf.py

"""
Modelo CRF para NER de Endereços - Compatível com sklearn Pipeline
"""
import sklearn_crfsuite
import re
from sklearn.base import BaseEstimator, TransformerMixin



class AddressFeatureExtractor(BaseEstimator, TransformerMixin):
    """Extrator de features morfológicas + contextuais para CRF (compatível com sklearn)"""

    def __init__(self):
        pass

    def fit(self, X, y=None):
        """Fit não faz nada, mas é necessário para compatibilidade sklearn

        Args:
            X: Lista de listas de tokens
            y: Ignorado (para compatibilidade sklearn)
        """
        return self

    def transform(self, X):
        """Transforma tokens em features para CRF

        Args:
            X: Lista de listas de tokens

        Returns:
            Lista de listas de dicionários de features
        """
        return [self._sent2features(tokens) for tokens in X]

    def _sent2features(self, tokens):
        """Converte sequência de tokens em features para CRF"""
        return [self._word2features(tokens, i) for i in range(len(tokens))]

    def _word2features(self, tokens, i):
        """Extrai features para CRF com contexto de vizinhança"""
        word = tokens[i]
        features = self._extract_word_features_crf(word)
        features['bias'] = 1.0

        if i > 0:
            word_prev = tokens[i - 1]
            features.update({
                '-1:word.lower': word_prev.lower(),
                '-1:is_title': word_prev.istitle(),
                '-1:is_upper': word_prev.isupper(),
                '-1:is_digit': word_prev.isdigit(),
            })
        else:
            features['BOS'] = True

        if i < len(tokens) - 1:
            word_next = tokens[i + 1]
            features.update({
                '+1:word.lower': word_next.lower(),
                '+1:is_title': word_next.istitle(),
                '+1:is_upper': word_next.isupper(),
                '+1:is_digit': word_next.isdigit(),
            })
        else:
            features['EOS'] = True

        return features

    def _extract_word_features_crf(self, word):
        """Extrai as MESMAS features morfológicas do HMM para o CRF"""
        features = {}

        word_lower = word.lower()
        features['word.lower'] = word_lower

        if word.isdigit():
            features['is_digit'] = True
            if len(word) == 1:
                features['digit_single'] = True
            elif len(word) <= 4:
                features['digit_small'] = True
            elif len(word) == 5:
                features['digit_five'] = True
            else:
                features['digit_large'] = True
        else:
            features['is_digit'] = False

        if word.isupper() and len(word) == 2 and word.isalpha():
            features['is_state_abbrev'] = True
        else:
            features['is_state_abbrev'] = False

        features['is_title'] = word.istitle()

        if word.isupper() and len(word) > 2:
            features['is_upper'] = True
        else:
            features['is_upper'] = False

        if '-' in word:
            features['has_dash'] = True
            if re.match(r'^\d{5}-\d{3}$', word):
                features['is_cep_format'] = True
            else:
                features['is_cep_format'] = False
        else:
            features['has_dash'] = False
            features['is_cep_format'] = False

        if re.match(r'^\d{8}$', word):
            features['is_cep_no_dash'] = True
        else:
            features['is_cep_no_dash'] = False

        word_len = len(word)
        if word_len == 1:
            features['len_1'] = True
        elif word_len == 2:
            features['len_2'] = True
        elif word_len <= 4:
            features['len_short'] = True
        elif word_len >= 10:
            features['len_long'] = True
        else:
            features['len_medium'] = True

        common_types = {
            'rua': 'tipo_rua', 'avenida': 'tipo_av', 'travessa': 'tipo_trav',
            'alameda': 'tipo_alam', 'praça': 'tipo_praca', 'praca': 'tipo_praca',
            'estrada': 'tipo_estrada', 'rodovia': 'tipo_rod', 'viela': 'tipo_viela',
            'av': 'tipo_av', 'r': 'tipo_rua', 'trav': 'tipo_trav'
        }
        if word_lower in common_types:
            features[common_types[word_lower]] = True

        complements = {'apto', 'ap', 'apartamento', 'sala', 'bloco', 'casa', 'lote', 'quadra'}
        if word_lower in complements:
            features['is_complement_word'] = True
        else:
            features['is_complement_word'] = False

        return features



class CRFTagger(BaseEstimator):
    """CRF para tagging de sequências (compatível com sklearn Pipeline)"""

    def __init__(self, algorithm='lbfgs', c1=0.1, c2=0.1, max_iterations=100,
                 all_possible_transitions=True):
        self.algorithm = algorithm
        self.c1 = c1
        self.c2 = c2
        self.max_iterations = max_iterations
        self.all_possible_transitions = all_possible_transitions

        self.model = None
        self.classes_ = None

    def fit(self, X, y):
        """Treina o modelo CRF

        Args:
            X: Lista de listas de dicionários de features
            y: Lista de listas de labels
        """
        self.model = sklearn_crfsuite.CRF(
            algorithm=self.algorithm,
            c1=self.c1,
            c2=self.c2,
            max_iterations=self.max_iterations,
            all_possible_transitions=self.all_possible_transitions,
            verbose=False
        )

        self.model.fit(X, y)

        self.classes_ = self.model.classes_

        print(f"  ✓ CRF treinado com {len(self.classes_)} labels")
        return self

    def predict(self, X):
        """Prediz labels para sequências

        Args:
            X: Lista de listas de dicionários de features

        Returns:
            Lista de listas de labels preditos
        """
        if self.model is None:
            raise ValueError("Modelo não treinado! Execute .fit() primeiro.")

        return self.model.predict(X)

    def get_params(self, deep=True):
        """Retorna parâmetros do estimator (sklearn compatibility)"""
        return {
            'algorithm': self.algorithm,
            'c1': self.c1,
            'c2': self.c2,
            'max_iterations': self.max_iterations,
            'all_possible_transitions': self.all_possible_transitions
        }

    def set_params(self, **params):
        """Define parâmetros do estimator (sklearn compatibility)"""
        for key, value in params.items():
            setattr(self, key, value)
        return self