Conditional Random Fields

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174

from typing import List, Tuple, Dict
from dataclasses import dataclass
 
@dataclass
class Entity:
    """Represents an extracted entity."""
    entity_type: str
    start: int  # Token index
    end: int    # Token index (exclusive)
    text: str
 
def bio_to_entities(tokens: List[str], labels: List[str]) -> List[Entity]:
    """
    Convert BIO labels to entity spans.
    
    Args:
        tokens: List of tokens
        labels: List of BIO labels
    
    Returns:
        List of Entity objects
    """
    entities = []
    current_entity = None
    current_start = None
    
    for i, (token, label) in enumerate(zip(tokens, labels)):
        if label.startswith('B-'):
            # Close previous entity if exists
            if current_entity:
                entities.append(Entity(
                    entity_type=current_entity,
                    start=current_start,
                    end=i,
                    text=' '.join(tokens[current_start:i])
                ))
            # Start new entity
            current_entity = label[2:]  # Remove 'B-' prefix
            current_start = i
            
        elif label.startswith('I-'):
            # Continue current entity (if matching type)
            entity_type = label[2:]
            if current_entity != entity_type:
                # Type mismatch - close previous and start new
                if current_entity:
                    entities.append(Entity(
                        entity_type=current_entity,
                        start=current_start,
                        end=i,
                        text=' '.join(tokens[current_start:i])
                    ))
                current_entity = entity_type
                current_start = i
                
        else:  # 'O' label
            # Close current entity if exists
            if current_entity:
                entities.append(Entity(
                    entity_type=current_entity,
                    start=current_start,
                    end=i,
                    text=' '.join(tokens[current_start:i])
                ))
            current_entity = None
            current_start = None
    
    # Don't forget last entity
    if current_entity:
        entities.append(Entity(
            entity_type=current_entity,
            start=current_start,
            end=len(tokens),
            text=' '.join(tokens[current_start:])
        ))
    
    return entities
 
 
def bioes_to_entities(tokens: List[str], labels: List[str]) -> List[Entity]:
    """
    Convert BIOES labels to entity spans.
    """
    entities = []
    current_entity = None
    current_start = None
    
    for i, (token, label) in enumerate(zip(tokens, labels)):
        if label.startswith('S-'):
            # Singleton entity
            entity_type = label[2:]
            entities.append(Entity(
                entity_type=entity_type,
                start=i,
                end=i+1,
                text=token
            ))
            current_entity = None
            
        elif label.startswith('B-'):
            # Start of multi-token entity
            current_entity = label[2:]
            current_start = i
            
        elif label.startswith('E-'):
            # End of multi-token entity
            if current_entity:
                entities.append(Entity(
                    entity_type=current_entity,
                    start=current_start,
                    end=i+1,
                    text=' '.join(tokens[current_start:i+1])
                ))
            current_entity = None
            
        elif label.startswith('I-'):
            # Inside - just continue
            pass
            
        else:  # 'O'
            current_entity = None
    
    return entities
 
 
def validate_bio_sequence(labels: List[str]) -> List[str]:
    """
    Validate and fix invalid BIO sequences.
    
    Invalid: I-X without preceding B-X or I-X
    Fix: Convert orphan I-X to B-X
    """
    fixed = []
    prev_entity = None
    
    for label in labels:
        if label.startswith('I-'):
            entity_type = label[2:]
            if prev_entity != entity_type:
                # Invalid - convert to B-
                fixed.append(f'B-{entity_type}')
                prev_entity = entity_type
            else:
                fixed.append(label)
        elif label.startswith('B-'):
            prev_entity = label[2:]
            fixed.append(label)
        else:
            prev_entity = None
            fixed.append(label)
    
    return fixed
 
 
# Example
tokens = ["Barack", "Obama", "was", "born", "in", "Honolulu", ",", "Hawaii"]
bio_labels = ["B-PER", "I-PER", "O", "O", "O", "B-LOC", "O", "B-LOC"]
bioes_labels = ["B-PER", "E-PER", "O", "O", "O", "S-LOC", "O", "S-LOC"]
 
print("BIO Example:")
for tok, label in zip(tokens, bio_labels):
    print(f"  {tok:12} {label}")
 
print("\nExtracted entities from BIO:")
for entity in bio_to_entities(tokens, bio_labels):
    print(f"  {entity.entity_type}: '{entity.text}' [{entity.start}:{entity.end}]")
 
print("\nBIOES Example:")
for tok, label in zip(tokens, bioes_labels):
    print(f"  {tok:12} {label}")
    
print("\nExtracted entities from BIOES:")
for entity in bioes_to_entities(tokens, bioes_labels):
    print(f"  {entity.entity_type}: '{entity.text}' [{entity.start}:{entity.end}]")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

from typing import List, Dict, Set
import re
 
class NERFeatureExtractor:
    """
    Feature extractor specialized for Named Entity Recognition.
    
    Implements standard NER features plus domain-specific patterns.
    """
    
    def __init__(
        self,
        person_gazetteer: Set[str] = None,
        org_gazetteer: Set[str] = None,
        loc_gazetteer: Set[str] = None,
    ):
        self.person_gaz = person_gazetteer or set()
        self.org_gaz = org_gazetteer or set()
        self.loc_gaz = loc_gazetteer or set()
    
    def extract_features(
        self, 
        tokens: List[str], 
        position: int
    ) -> Dict[str, float]:
        """
        Extract features for NER at given position.
        
        Returns dictionary of feature_name -> value.
        """
        features = {}
        n = len(tokens)
        word = tokens[position]
        
        # ============== Word Identity ==============
        features['word.lower()'] = word.lower()
        
        # Previous/next words
        if position > 0:
            features['word[-1].lower()'] = tokens[position-1].lower()
        else:
            features['BOS'] = True
            
        if position < n - 1:
            features['word[+1].lower()'] = tokens[position+1].lower()
        else:
            features['EOS'] = True
        
        # ============== Orthographic ==============
        features['is_capitalized'] = word[0].isupper() if word else False
        features['is_all_caps'] = word.isupper()
        features['is_all_lower'] = word.islower()
        features['is_title'] = word.istitle()
        
        # Mixed case patterns
        features['has_mixed_case'] = (
            any(c.isupper() for c in word) and 
            any(c.islower() for c in word) and
            not word.istitle()
        )
        
        # ============== Shape ==============
        features['word_shape'] = self._get_word_shape(word)
        features['short_word_shape'] = self._get_short_shape(word)
        
        # ============== Affixes ==============
        for length in [2, 3, 4]:
            if len(word) >= length:
                features[f'prefix_{length}'] = word[:length].lower()
                features[f'suffix_{length}'] = word[-length:].lower()
        
        # ============== Character Patterns ==============
        features['has_digit'] = any(c.isdigit() for c in word)
        features['all_digits'] = word.isdigit()
        features['has_hyphen'] = '-' in word
        features['has_period'] = '.' in word
        features['has_apostrophe'] = "'" in word
        
        # Number patterns
        features['is_year'] = bool(re.match(r'^(19|20)\d{2}$', word))
        features['is_money'] = bool(re.match(r'^\$[\d,]+(\.\d{2})?$', word))
        features['is_percent'] = bool(re.match(r'^\d+(\.\d+)?%$', word))
        
        # ============== Gazetteer Lookups ==============
        word_lower = word.lower()
        features['in_person_gaz'] = word_lower in self.person_gaz
        features['in_org_gaz'] = word_lower in self.org_gaz
        features['in_loc_gaz'] = word_lower in self.loc_gaz
        
        # ============== Context Features ==============
        # Position in sentence
        features['is_first_word'] = (position == 0)
        features['is_last_word'] = (position == n - 1)
        
        # Surrounding capitalization
        if position > 0:
            features['prev_is_cap'] = tokens[position-1][0].isupper() if tokens[position-1] else False
        if position < n - 1:
            features['next_is_cap'] = tokens[position+1][0].isupper() if tokens[position+1] else False
        
        # ============== Trigger Words ==============
        # Words that often precede entities
        if position > 0:
            prev = tokens[position-1].lower()
            features['prev_is_title'] = prev in {'mr', 'mrs', 'ms', 'dr', 'prof', 'president', 'ceo'}
            features['prev_is_prep'] = prev in {'in', 'at', 'from', 'to', 'of'}
        
        return features
    
    def _get_word_shape(self, word: str) -> str:
        """Convert word to shape pattern."""
        shape = []
        for c in word:
            if c.isupper():
                shape.append('X')
            elif c.islower():
                shape.append('x')
            elif c.isdigit():
                shape.append('d')
            else:
                shape.append(c)
        return ''.join(shape)
    
    def _get_short_shape(self, word: str) -> str:
        """Collapsed shape pattern."""
        shape = self._get_word_shape(word)
        result = [shape[0]] if shape else []
        for c in shape[1:]:
            if c != result[-1]:
                result.append(c)
        return ''.join(result)
 
 
# Example
tokens = ["President", "Barack", "Obama", "visited", "the", "United", "Nations", "in", "New", "York"]
 
extractor = NERFeatureExtractor(
    person_gazetteer={'obama', 'biden', 'trump'},
    org_gazetteer={'united nations', 'google', 'apple'},
    loc_gazetteer={'new york', 'paris', 'london'}
)
 
print("NER Features for 'Barack' (position 1):")
features = extractor.extract_features(tokens, 1)
for name, value in sorted(features.items())[:15]:
    print(f"  {name}: {value}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

from typing import List, Dict
 
class POSFeatureExtractor:
    """
    Feature extractor for Part-of-Speech tagging.
    """
    
    def __init__(self):
        self.common_suffixes = {
            'verb': ['ing', 'ed', 'en', 'es', 's'],
            'noun': ['tion', 'ness', 'ment', 'ity', 'er', 'or'],
            'adj': ['al', 'ful', 'less', 'ous', 'ive', 'able', 'ible'],
            'adv': ['ly']
        }
    
    def extract_features(
        self, 
        tokens: List[str], 
        position: int
    ) -> Dict[str, float]:
        """
        Extract POS-specific features.
        """
        features = {}
        n = len(tokens)
        word = tokens[position]
        word_lower = word.lower()
        
        # ============== Word Features ==============
        features['word'] = word_lower
        
        # Window words
        for offset in [-2, -1, 1, 2]:
            idx = position + offset
            if 0 <= idx < n:
                features[f'word[{offset:+d}]'] = tokens[idx].lower()
        
        # ============== Morphological ==============
        # Suffixes (crucial for POS)
        for length in [1, 2, 3, 4]:
            if len(word) >= length:
                suffix = word_lower[-length:]
                features[f'suffix{length}'] = suffix
                
                # Check against known POS suffixes
                for pos, suffixes in self.common_suffixes.items():
                    if suffix in suffixes:
                        features[f'has_{pos}_suffix'] = True
        
        # Prefixes
        for length in [2, 3]:
            if len(word) >= length:
                features[f'prefix{length}'] = word_lower[:length]
        
        # ============== Orthographic ==============
        features['is_cap'] = word[0].isupper() if word else False
        features['all_caps'] = word.isupper() and len(word) > 1
        features['has_digit'] = any(c.isdigit() for c in word)
        features['has_hyphen'] = '-' in word
        features['all_digits'] = word.isdigit()
        
        # ============== Word Patterns ==============
        features['is_punct'] = word in '.,;:!?()[]{}"\'/-'
        features['is_quote'] = word in '"\'``\'\''
        
        # Length features
        features['word_length'] = len(word)
        features['is_short'] = len(word) <= 2
        
        # ============== Context Patterns ==============
        # Bigram features
        if position > 0:
            features['bigram[-1,0]'] = f"{tokens[position-1].lower()}_{word_lower}"
        if position < n - 1:
            features['bigram[0,+1]'] = f"{word_lower}_{tokens[position+1].lower()}"
        
        # ============== Position ==============
        features['is_first'] = (position == 0)
        features['is_last'] = (position == n - 1)
        
        return features
 
 
# Example
sentence = "The quick brown fox jumps over the lazy dog".split()
 
extractor = POSFeatureExtractor()
print("POS Features for 'jumps' (position 4):")
features = extractor.extract_features(sentence, 4)
for name, value in sorted(features.items())[:12]:
    print(f"  {name}: {value}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

from typing import List, Dict, Set, Tuple
from collections import defaultdict
from dataclasses import dataclass
 
@dataclass
class Entity:
    entity_type: str
    start: int
    end: int
    text: str = ""
    
    def __hash__(self):
        return hash((self.entity_type, self.start, self.end))
    
    def __eq__(self, other):
        return (self.entity_type == other.entity_type and 
                self.start == other.start and 
                self.end == other.end)
 
 
def compute_ner_metrics(
    gold_entities: List[Entity],
    pred_entities: List[Entity]
) -> Dict[str, float]:
    """
    Compute NER evaluation metrics.
    
    Uses exact match: both boundaries and type must match.
    """
    gold_set = set(gold_entities)
    pred_set = set(pred_entities)
    
    true_positives = gold_set & pred_set
    false_positives = pred_set - gold_set
    false_negatives = gold_set - pred_set
    
    tp = len(true_positives)
    fp = len(false_positives)
    fn = len(false_negatives)
    
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
    
    return {
        'precision': precision,
        'recall': recall,
        'f1': f1,
        'true_positives': tp,
        'false_positives': fp,
        'false_negatives': fn
    }
 
 
def compute_metrics_by_type(
    gold_entities: List[Entity],
    pred_entities: List[Entity]
) -> Dict[str, Dict[str, float]]:
    """
    Compute metrics separately for each entity type.
    """
    # Group by type
    gold_by_type = defaultdict(list)
    pred_by_type = defaultdict(list)
    
    for entity in gold_entities:
        gold_by_type[entity.entity_type].append(entity)
    for entity in pred_entities:
        pred_by_type[entity.entity_type].append(entity)
    
    all_types = set(gold_by_type.keys()) | set(pred_by_type.keys())
    
    results = {}
    for entity_type in all_types:
        results[entity_type] = compute_ner_metrics(
            gold_by_type[entity_type],
            pred_by_type[entity_type]
        )
    
    return results
 
 
def compute_seqeval_metrics(
    gold_labels: List[List[str]],
    pred_labels: List[List[str]]
) -> Dict[str, float]:
    """
    Compute metrics using seqeval-style evaluation.
    
    This is the standard evaluation used in research papers.
    """
    all_gold_entities = []
    all_pred_entities = []
    
    for sent_idx, (gold, pred) in enumerate(zip(gold_labels, pred_labels)):
        # Extract entities from BIO labels
        gold_ents = extract_entities_with_offset(gold, sent_idx * 1000)
        pred_ents = extract_entities_with_offset(pred, sent_idx * 1000)
        
        all_gold_entities.extend(gold_ents)
        all_pred_entities.extend(pred_ents)
    
    return compute_ner_metrics(all_gold_entities, all_pred_entities)
 
 
def extract_entities_with_offset(labels: List[str], offset: int) -> List[Entity]:
    """Extract entities with position offset for unique identification."""
    entities = []
    current_type = None
    current_start = None
    
    for i, label in enumerate(labels):
        if label.startswith('B-'):
            if current_type:
                entities.append(Entity(current_type, current_start + offset, i + offset))
            current_type = label[2:]
            current_start = i
        elif label.startswith('I-'):
            if current_type is None or label[2:] != current_type:
                if current_type:
                    entities.append(Entity(current_type, current_start + offset, i + offset))
                current_type = label[2:]
                current_start = i
        else:
            if current_type:
                entities.append(Entity(current_type, current_start + offset, i + offset))
            current_type = None
    
    if current_type:
        entities.append(Entity(current_type, current_start + offset, len(labels) + offset))
    
    return entities
 
 
# Example evaluation
gold = [Entity('PER', 0, 2), Entity('LOC', 5, 6), Entity('ORG', 8, 10)]
pred = [Entity('PER', 0, 2), Entity('LOC', 5, 7), Entity('ORG', 8, 10)]  # LOC boundary wrong
 
print("NER Evaluation Example")
print("=" * 50)
print(f"Gold entities: {[(e.entity_type, e.start, e.end) for e in gold]}")
print(f"Pred entities: {[(e.entity_type, e.start, e.end) for e in pred]}")
 
metrics = compute_ner_metrics(gold, pred)
print(f"\nMetrics:")
print(f"  Precision: {metrics['precision']:.3f}")
print(f"  Recall:    {metrics['recall']:.3f}")
print(f"  F1:        {metrics['f1']:.3f}")
print(f"  TP: {metrics['true_positives']}, FP: {metrics['false_positives']}, FN: {metrics['false_negatives']}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367

"""
Complete CRF Sequence Labeling Pipeline
 
This module implements a full NER system using CRFs:
1. Data loading and preprocessing
2. Feature extraction
3. Model training with regularization
4. Evaluation on test set
5. Prediction on new text
"""
 
import numpy as np
from typing import List, Tuple, Dict
from collections import defaultdict
import pickle
 
class CRFSequenceLabeler:
    """
    Complete CRF-based sequence labeler.
    """
    
    def __init__(
        self,
        feature_templates: List[str] = None,
        regularization: float = 1.0
    ):
        self.feature_templates = feature_templates or [
            'word', 'word_lower', 'suffix2', 'suffix3',
            'prefix2', 'prefix3', 'is_capitalized', 'word_shape'
        ]
        self.regularization = regularization
        
        self.feature_to_id = {}
        self.label_to_id = {}
        self.id_to_label = {}
        self.weights = None
    
    def fit(
        self,
        X: List[List[str]],  # List of token sequences
        y: List[List[str]],  # List of label sequences
        epochs: int = 10,
        learning_rate: float = 0.1,
        verbose: bool = True
    ):
        """
        Train the CRF model.
        
        Args:
            X: List of token sequences
            y: List of label sequences
            epochs: Number of training epochs
            learning_rate: SGD learning rate
            verbose: Print training progress
        """
        # Step 1: Build vocabularies
        self._build_vocabularies(X, y)
        
        if verbose:
            print(f"Vocabulary: {len(self.feature_to_id)} features, {len(self.label_to_id)} labels")
        
        # Step 2: Initialize weights
        self.weights = np.zeros(len(self.feature_to_id))
        
        # Step 3: Training loop (SGD)
        for epoch in range(epochs):
            total_loss = 0.0
            
            # Shuffle data
            indices = np.random.permutation(len(X))
            
            for idx in indices:
                tokens, labels = X[idx], y[idx]
                label_ids = [self.label_to_id[l] for l in labels]
                
                # Compute gradient and loss
                gradient, loss = self._compute_gradient(tokens, label_ids)
                
                # Add regularization gradient
                gradient -= self.regularization * self.weights
                reg_loss = 0.5 * self.regularization * np.sum(self.weights ** 2)
                
                # SGD update
                self.weights += learning_rate * gradient
                total_loss += loss + reg_loss
            
            if verbose:
                avg_loss = total_loss / len(X)
                print(f"Epoch {epoch+1}/{epochs}: loss = {avg_loss:.4f}")
    
    def predict(self, X: List[List[str]]) -> List[List[str]]:
        """
        Predict labels for sequences.
        
        Args:
            X: List of token sequences
        
        Returns:
            List of predicted label sequences
        """
        predictions = []
        
        for tokens in X:
            label_ids = self._viterbi(tokens)
            labels = [self.id_to_label[lid] for lid in label_ids]
            predictions.append(labels)
        
        return predictions
    
    def _build_vocabularies(self, X: List[List[str]], y: List[List[str]]):
        """Build feature and label vocabularies from training data."""
        # Collect labels
        labels = set()
        for label_seq in y:
            labels.update(label_seq)
        
        self.label_to_id = {l: i for i, l in enumerate(sorted(labels))}
        self.id_to_label = {i: l for l, i in self.label_to_id.items()}
        
        # Collect features (feature_name, label combinations)
        feature_set = set()
        L = len(self.label_to_id)
        
        for tokens in X:
            for i in range(len(tokens)):
                feat_dict = self._extract_features(tokens, i)
                for feat_name in feat_dict:
                    for label_id in range(L):
                        feature_set.add((feat_name, label_id))
        
        # Add transition features
        for l1 in range(L + 1):  # +1 for START
            for l2 in range(L):
                feature_set.add((f'trans_{l1}_{l2}', -1))  # -1 indicates transition
        
        self.feature_to_id = {f: i for i, f in enumerate(sorted(feature_set))}
    
    def _extract_features(self, tokens: List[str], position: int) -> Dict[str, float]:
        """Extract features at a position."""
        features = {}
        word = tokens[position]
        n = len(tokens)
        
        features['word'] = word
        features['word_lower'] = word.lower()
        features['is_cap'] = word[0].isupper() if word else False
        
        # Suffixes and prefixes
        for length in [2, 3]:
            if len(word) >= length:
                features[f'suffix{length}'] = word[-length:].lower()
                features[f'prefix{length}'] = word[:length].lower()
        
        # Context words
        if position > 0:
            features['prev_word'] = tokens[position-1].lower()
        else:
            features['BOS'] = True
            
        if position < n - 1:
            features['next_word'] = tokens[position+1].lower()
        else:
            features['EOS'] = True
        
        return features
    
    def _compute_gradient(
        self, 
        tokens: List[str], 
        label_ids: List[int]
    ) -> Tuple[np.ndarray, float]:
        """Compute gradient for a single sequence."""
        n = len(tokens)
        L = len(self.label_to_id)
        
        # Run forward-backward
        log_alpha, log_Z = self._forward(tokens)
        log_beta = self._backward(tokens)
        
        gradient = np.zeros(len(self.feature_to_id))
        
        # Observed features (correct labels)
        for i in range(n):
            y = label_ids[i]
            y_prev = label_ids[i-1] if i > 0 else L  # L = START
            
            # Emission features
            feat_dict = self._extract_features(tokens, i)
            for feat_name, value in feat_dict.items():
                key = (feat_name, y)
                if key in self.feature_to_id:
                    gradient[self.feature_to_id[key]] += 1.0
            
            # Transition features
            trans_key = (f'trans_{y_prev}_{y}', -1)
            if trans_key in self.feature_to_id:
                gradient[self.feature_to_id[trans_key]] += 1.0
        
        # Expected features (under model distribution)
        for i in range(n):
            feat_dict = self._extract_features(tokens, i)
            
            for y in range(L):
                # Marginal P(y_i = y | x)
                log_marginal = log_alpha[i, y] + log_beta[i, y] - log_Z
                marginal = np.exp(log_marginal)
                
                for feat_name in feat_dict:
                    key = (feat_name, y)
                    if key in self.feature_to_id:
                        gradient[self.feature_to_id[key]] -= marginal
            
            # Transition expected features
            if i == 0:
                for y in range(L):
                    trans_key = (f'trans_{L}_{y}', -1)  # START -> y
                    if trans_key in self.feature_to_id:
                        log_marg = log_alpha[0, y] + log_beta[0, y] - log_Z
                        gradient[self.feature_to_id[trans_key]] -= np.exp(log_marg)
            else:
                for y_prev in range(L):
                    for y in range(L):
                        trans_key = (f'trans_{y_prev}_{y}', -1)
                        if trans_key in self.feature_to_id:
                            # Pairwise marginal
                            trans_score = self._get_transition_score(y_prev, y)
                            emit_score = self._get_emission_score(tokens, i, y)
                            log_marg = (log_alpha[i-1, y_prev] + trans_score + 
                                       emit_score + log_beta[i, y] - log_Z)
                            gradient[self.feature_to_id[trans_key]] -= np.exp(log_marg)
        
        # Compute loss = -log P(y|x)
        score = sum(
            self._get_emission_score(tokens, i, label_ids[i]) +
            self._get_transition_score(label_ids[i-1] if i > 0 else L, label_ids[i])
            for i in range(n)
        )
        loss = log_Z - score
        
        return gradient, loss
    
    def _forward(self, tokens: List[str]) -> Tuple[np.ndarray, float]:
        """Forward algorithm."""
        n = len(tokens)
        L = len(self.label_to_id)
        
        log_alpha = np.full((n, L), -np.inf)
        
        # Initialize
        for y in range(L):
            trans = self._get_transition_score(L, y)  # START -> y
            emit = self._get_emission_score(tokens, 0, y)
            log_alpha[0, y] = trans + emit
        
        # Recursion
        for i in range(1, n):
            for y in range(L):
                emit = self._get_emission_score(tokens, i, y)
                scores = [
                    log_alpha[i-1, y_prev] + self._get_transition_score(y_prev, y) + emit
                    for y_prev in range(L)
                ]
                log_alpha[i, y] = self._logsumexp(np.array(scores))
        
        log_Z = self._logsumexp(log_alpha[n-1])
        return log_alpha, log_Z
    
    def _backward(self, tokens: List[str]) -> np.ndarray:
        """Backward algorithm."""
        n = len(tokens)
        L = len(self.label_to_id)
        
        log_beta = np.full((n, L), -np.inf)
        log_beta[n-1] = 0.0
        
        for i in range(n-2, -1, -1):
            for y in range(L):
                scores = [
                    self._get_transition_score(y, y_next) +
                    self._get_emission_score(tokens, i+1, y_next) +
                    log_beta[i+1, y_next]
                    for y_next in range(L)
                ]
                log_beta[i, y] = self._logsumexp(np.array(scores))
        
        return log_beta
    
    def _viterbi(self, tokens: List[str]) -> List[int]:
        """Viterbi decoding."""
        n = len(tokens)
        L = len(self.label_to_id)
        
        delta = np.full((n, L), -np.inf)
        psi = np.zeros((n, L), dtype=int)
        
        # Initialize
        for y in range(L):
            delta[0, y] = (self._get_transition_score(L, y) + 
                          self._get_emission_score(tokens, 0, y))
        
        # Recursion
        for i in range(1, n):
            for y in range(L):
                emit = self._get_emission_score(tokens, i, y)
                scores = [
                    delta[i-1, y_prev] + self._get_transition_score(y_prev, y) + emit
                    for y_prev in range(L)
                ]
                psi[i, y] = np.argmax(scores)
                delta[i, y] = scores[psi[i, y]]
        
        # Backtrack
        path = [np.argmax(delta[n-1])]
        for i in range(n-1, 0, -1):
            path.append(psi[i, path[-1]])
        
        return list(reversed(path))
    
    def _get_emission_score(self, tokens: List[str], pos: int, label: int) -> float:
        """Get emission score for position and label."""
        feat_dict = self._extract_features(tokens, pos)
        score = 0.0
        for feat_name in feat_dict:
            key = (feat_name, label)
            if key in self.feature_to_id:
                score += self.weights[self.feature_to_id[key]]
        return score
    
    def _get_transition_score(self, prev_label: int, label: int) -> float:
        """Get transition score."""
        key = (f'trans_{prev_label}_{label}', -1)
        if key in self.feature_to_id:
            return self.weights[self.feature_to_id[key]]
        return 0.0
    
    @staticmethod
    def _logsumexp(arr: np.ndarray) -> float:
        max_val = np.max(arr)
        if max_val == -np.inf:
            return -np.inf
        return max_val + np.log(np.sum(np.exp(arr - max_val)))
 
 
# Example usage
if __name__ == "__main__":
    # Sample training data
    X_train = [
        ["John", "works", "at", "Google", "."],
        ["Mary", "lives", "in", "Paris", "."],
    ]
    y_train = [
        ["B-PER", "O", "O", "B-ORG", "O"],
        ["B-PER", "O", "O", "B-LOC", "O"],
    ]
    
    # Train model
    model = CRFSequenceLabeler(regularization=0.1)
    model.fit(X_train, y_train, epochs=50, learning_rate=0.1, verbose=True)
    
    # Predict
    X_test = [["Tom", "joined", "Microsoft", "."]]
    predictions = model.predict(X_test)
    
    print("\nPredictions:")
    for tokens, labels in zip(X_test, predictions):
        for tok, lab in zip(tokens, labels):
            print(f"  {tok}: {lab}")