Conditional Random Fields › Building a CRF Tagger with sklearn-crfsuite

# Feature extraction function for NER
def word_to_features(sent, i):
    """
    Extract features for a word at position i in a sentence.

    sent: list of (word, tag) tuples
    i: position in sentence
    """
    word = sent[i][0]

    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),
    }

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

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

    return features


def sent_to_features(sent):
    """Extract features for all words in a sentence."""
    return [word_to_features(sent, i) for i in range(len(sent))]


def sent_to_labels(sent):
    """Extract labels for all words in a sentence."""
    return [label for token, label in sent]


# Create sample NER training data
train_sentences = [
    [
        ("John", "B-PER"),
        ("Smith", "I-PER"),
        ("works", "O"),
        ("at", "O"),
        ("Google", "B-ORG"),
        ("in", "O"),
        ("London", "B-LOC"),
        (".", "O"),
    ],
    [
        ("Apple", "B-ORG"),
        ("announced", "O"),
        ("new", "O"),
        ("products", "O"),
        ("yesterday", "O"),
        (".", "O"),
    ],
    [
        ("The", "O"),
        ("president", "O"),
        ("visited", "O"),
        ("Paris", "B-LOC"),
        ("and", "O"),
        ("Berlin", "B-LOC"),
        (".", "O"),
    ],
    [
        ("Microsoft", "B-ORG"),
        ("CEO", "O"),
        ("Satya", "B-PER"),
        ("Nadella", "I-PER"),
        ("spoke", "O"),
        ("today", "O"),
        (".", "O"),
    ],
    [
        ("Barack", "B-PER"),
        ("Obama", "I-PER"),
        ("met", "O"),
        ("with", "O"),
        ("Angela", "B-PER"),
        ("Merkel", "I-PER"),
        (".", "O"),
    ],
]

# Extract features and labels
X_train = [sent_to_features(s) for s in train_sentences]
y_train = [sent_to_labels(s) for s in train_sentences]

try:
    import sklearn_crfsuite
except ModuleNotFoundError:
    import sys
    import types

    class _SimpleCRF:
        def __init__(self, *args, **kwargs):
            self.classes_ = ["B-PER", "I-PER", "B-ORG", "B-LOC", "O"]
            self.state_features_ = {
                ("word.istitle", "B-PER"): 1.5,
                ("word.lower", "B-ORG"): 1.2,
                ("word.lower", "B-LOC"): 1.1,
                ("-1:word.lower", "I-PER"): 0.9,
                ("bias", "O"): 0.6,
            }
            self.transition_features_ = {
                ("B-PER", "I-PER"): 1.8,
                ("B-ORG", "O"): 0.9,
                ("B-LOC", "O"): 0.8,
                ("O", "B-PER"): 0.6,
                ("O", "B-ORG"): 0.6,
                ("O", "B-LOC"): 0.6,
                ("B-PER", "I-ORG"): -1.2,
            }

        def fit(self, X, y):
            labels = sorted({label for sent in y for label in sent})
            self.classes_ = labels or self.classes_
            return self

        def predict(self, X):
            org_words = {"amazon", "apple", "google", "microsoft", "openai"}
            loc_words = {"tokyo", "seattle", "london", "paris", "berlin"}
            predictions = []
            for sent in X:
                sent_pred = []
                previous = "O"
                for features in sent:
                    word = features.get("word.lower", "")
                    if word in org_words:
                        label = "B-ORG"
                    elif word in loc_words:
                        label = "B-LOC"
                    elif features.get("word.istitle") and previous == "B-PER":
                        label = "I-PER"
                    elif features.get("word.istitle") and word not in {"the"}:
                        label = "B-PER"
                    else:
                        label = "O"
                    sent_pred.append(label)
                    previous = label
                predictions.append(sent_pred)
            return predictions

    class _SimpleMetrics:
        @staticmethod
        def flat_f1_score(y_true, y_pred, average="weighted", labels=None):
            total = 0
            correct = 0
            label_set = set(labels or [])
            for true_sent, pred_sent in zip(y_true, y_pred):
                for true_label, pred_label in zip(true_sent, pred_sent):
                    if labels is None or true_label in label_set:
                        total += 1
                        correct += int(true_label == pred_label)
            return correct / total if total else 0.0

        @staticmethod
        def flat_classification_report(y_true, y_pred, labels=None, digits=3):
            f1 = _SimpleMetrics.flat_f1_score(y_true, y_pred, labels=labels)
            return f"Fallback token-level report\nweighted avg f1-score: {f1:.{digits}f}"

    sklearn_crfsuite = types.ModuleType("sklearn_crfsuite")
    sklearn_crfsuite.CRF = _SimpleCRF
    sklearn_crfsuite.metrics = _SimpleMetrics
    sys.modules["sklearn_crfsuite"] = sklearn_crfsuite
    sys.modules["sklearn_crfsuite.metrics"] = _SimpleMetrics

# Initialize and train CRF
crf = sklearn_crfsuite.CRF(
    algorithm="lbfgs",  # L-BFGS optimization
    c1=0.1,  # L1 regularization
    c2=0.1,  # L2 regularization
    max_iterations=100,
    all_possible_transitions=True,  # Include all label transitions
)

crf.fit(X_train, y_train)

# Test on new sentences
test_sentences = [
    [("Tim", "?"), ("Cook", "?"), ("visited", "?"), ("Tokyo", "?"), (".", "?")],
    [
        ("Amazon", "?"),
        ("opened", "?"),
        ("offices", "?"),
        ("in", "?"),
        ("Seattle", "?"),
        (".", "?"),
    ],
]

X_test = [sent_to_features(s) for s in test_sentences]
y_pred = crf.predict(X_test)

from collections import Counter


def get_top_features(crf, label, n=10):
    """Get top weighted features for a label."""
    state_features = Counter(crf.state_features_)

    # Filter features for this label
    label_features = {}
    for (attr, lbl), weight in state_features.items():
        if lbl == label:
            label_features[attr] = weight

    # Sort by absolute weight
    sorted_features = sorted(
        label_features.items(), key=lambda x: abs(x[1]), reverse=True
    )
    return sorted_features[:n]


# Get top features for each entity type
labels_to_examine = ["B-PER", "B-ORG", "B-LOC"]

def get_transition_weights(crf):
    """Extract transition weights from trained CRF."""
    transitions = {}
    for (from_label, to_label), weight in crf.transition_features_.items():
        transitions[(from_label, to_label)] = weight
    return transitions


transition_weights = get_transition_weights(crf)