Conditional Random Fields › CRF for Named Entity Recognition

from nltk.corpus import conll2002

# Download the corpus if needed (run once locally, then comment out)
# nltk.download("conll2002", quiet=True)

# Load Spanish NER data (smaller than English, good for demonstration)
train_sents = list(conll2002.iob_sents("esp.train"))[
    :500
]  # Use subset for speed
test_sents = list(conll2002.iob_sents("esp.testa"))[:100]

def word_to_ner_features(sent, i):
    """
    Extract comprehensive features for NER.

    sent: list of (word, pos, ner) tuples
    i: position
    """
    word = sent[i][0]
    pos = sent[i][1]

    features = {
        "bias": 1.0,
        "word.lower": word.lower(),
        "word[-3:]": word[-3:],
        "word[-2:]": word[-2:],
        "word[:3]": word[:3],
        "word[:2]": word[:2],
        "word.isupper": word.isupper(),
        "word.istitle": word.istitle(),
        "word.isdigit": word.isdigit(),
        "word.length": len(word),
        "postag": pos,
        "postag[:2]": pos[:2],
    }

    # Previous word features
    if i > 0:
        word_prev = sent[i - 1][0]
        pos_prev = sent[i - 1][1]
        features.update(
            {
                "-1:word.lower": word_prev.lower(),
                "-1:word.istitle": word_prev.istitle(),
                "-1:word.isupper": word_prev.isupper(),
                "-1:postag": pos_prev,
                "-1:postag[:2]": pos_prev[:2],
            }
        )
    else:
        features["BOS"] = True

    # Next word features
    if i < len(sent) - 1:
        word_next = sent[i + 1][0]
        pos_next = sent[i + 1][1]
        features.update(
            {
                "+1:word.lower": word_next.lower(),
                "+1:word.istitle": word_next.istitle(),
                "+1:word.isupper": word_next.isupper(),
                "+1:postag": pos_next,
                "+1:postag[:2]": pos_next[:2],
            }
        )
    else:
        features["EOS"] = True

    return features


def sent_to_ner_features(sent):
    return [word_to_ner_features(sent, i) for i in range(len(sent))]


def sent_to_ner_labels(sent):
    return [ner for word, pos, ner in sent]


# Prepare data
X_train_ner = [sent_to_ner_features(s) for s in train_sents]
y_train_ner = [sent_to_ner_labels(s) for s in train_sents]
X_test_ner = [sent_to_ner_features(s) for s in test_sents]
y_test_ner = [sent_to_ner_labels(s) for s in test_sents]

# Train CRF
crf_ner = sklearn_crfsuite.CRF(
    algorithm="lbfgs",
    c1=0.1,
    c2=0.1,
    max_iterations=100,
    all_possible_transitions=True,
)

crf_ner.fit(X_train_ner, y_train_ner)

# Predict
y_pred_ner = crf_ner.predict(X_test_ner)

# Get labels (excluding O for focused evaluation)
labels = list(crf_ner.classes_)
labels.remove("O")