anno/backends/crf/

mod.rs

//! Conditional Random Field (CRF) NER backend.
//!
//! Implements classical statistical NER using CRF sequence labeling.
//! This provides a lightweight, interpretable baseline that requires no
//! external dependencies or GPU acceleration.
//!
//! # History
//!
//! CRF-based NER was a common baseline throughout the 2000s (pre-neural sequence labeling):
//! - Lafferty et al. 2001: Introduced CRFs for sequence labeling (ICML)
//! - McCallum & Li 2003: Applied CRFs to NER
//! - Stanford NER (2003-2014): CRF-based, still widely used
//! - State-of-art until neural methods (BiLSTM-CRF, 2015+)
//!
//! # Why CRF Beat Previous Methods
//!
//! CRFs solved the **label bias problem** that plagued MEMMs (Maximum Entropy
//! Markov Models, McCallum et al. 2000):
//!
//! ```text
//! Label Bias: In MEMMs, states with few successors effectively ignore
//! observations. Transition scores are conditional on current state,
//! so low-entropy states "absorb" probability mass regardless of input.
//!
//! HMM:   Generative model     P(x,y) = P(y) × P(x|y)
//! MEMM:  Local discriminative P(y_t|y_{t-1}, x)  ← label bias here
//! CRF:   Global discriminative P(y|x) = (1/Z) exp(∑ features × weights)
//!                                       ↑ normalizes over entire sequence
//! ```
//!
//! CRF models the conditional probability of the entire label sequence given
//! the observation sequence, using global normalization:
//!
//! ```text
//! P(y|x) = (1/Z(x)) × exp( ∑_t ∑_k λ_k × f_k(y_t, y_{t-1}, x, t) )
//!
//! where:
//!   - Z(x) is the partition function (normalizer)
//!   - f_k are feature functions
//!   - λ_k are learned weights
//! ```
//!
//! # References
//!
//! - Lafferty, McCallum, Pereira (2001): "Conditional Random Fields:
//!   Probabilistic Models for Segmenting and Labeling Sequence Data" (ICML)
//! - McCallum & Li (2003): "Early Results for Named Entity Recognition with
//!   Conditional Random Fields, Feature Induction and Web-Enhanced Lexicons"
//! - Finkel, Grenager, Manning (2005): "Incorporating Non-local Information
//!   into Information Extraction Systems by Gibbs Sampling" (ACL)
//!
//! # See Also
//!
//! - Historical NER baselines (HMM/CRF-era sequence models)
//!
//! # Feature Templates
//!
//! The CRF uses the following feature templates (matching `train_crf_weights.py`):
//!
//! ```text
//! - bias           : Always-on feature for label-specific bias
//! - word.lower     : Lowercased current word
//! - word.shape     : Word shape pattern (Xx, X, x, 0, etc.)
//! - word.isdigit   : Whether word is all digits
//! - word.istitle   : Whether word is titlecased
//! - word.isupper   : Whether word is all uppercase
//! - prefix{2,3}    : First 2-3 characters
//! - suffix{2,3}    : Last 2-3 characters
//! - -1:word.*      : Previous word features
//! - +1:word.*      : Next word features
//! - BOS/EOS        : Sentence boundary markers
//! ```
//!
//! # Performance
//!
//! Performance depends on weights, tokenization, and dataset; use the eval harness
//! for quantitative results.
//! | Heuristic | Lower | Hand-tuned, always available |
//! | Trained | Higher | From `train_crf_weights.py` |
//! | Neural | Highest | For comparison (GLiNER, BERT) |
//!
//! # Usage
//!
//! ```rust
//! use anno::CrfNER;
//! use anno::Model;
//!
//! // Use with default heuristic weights
//! let ner = CrfNER::new();
//! let entities = ner.extract_entities("John Smith works at Google", None)?;
//!
//! // Or load trained weights for better accuracy
//! // let ner = CrfNER::with_weights("crf_weights.json")?;
//! # Ok::<(), anno::Error>(())
//! ```
//!
//! # Training Weights
//!
//! To train weights on CoNLL-2003:
//!
//! ```bash
//! uv run scripts/train_crf_weights.py
//! ```
//!
//! This produces `crf_weights.json` which can be loaded with `CrfNER::with_weights()`.
//!
//! Nuance: CoNLL-2003’s English text is derived from Reuters/RCV1 and is commonly treated as
//! redistribution-restricted. The CoNLL site notes that, “because of copyright reasons we only
//! make available the annotations” and that you need separate access to the Reuters corpus to
//! build the full dataset: `http://www.clips.uantwerpen.be/conll2003/ner/`.
//!
//! Practical consequence: `anno` includes a training script, but it does not ship a CoNLL-trained
//! `crf_weights.json` out of the box.
//!
//! # Advantages Over Neural Methods
//!
//! - **Interpretable**: Features and weights are human-readable
//! - **Fast training**: CPU-only training; typically faster to iterate than neural training loops
//! - **No dependencies**: Pure Rust, no ONNX/Candle required
//! - **Deterministic**: Same input always produces same output
//! - **Small footprint**: Small weights file compared to ML model artifacts

use crate::{Entity, EntityType, Model, Result};
use std::collections::HashMap;
#[cfg(feature = "bundled-crf-weights")]
use std::sync::OnceLock;

/// CRF-based NER model.
///
/// Uses hand-crafted features and sequence labeling for named entity recognition.
/// This is a pure-Rust implementation that doesn't require external libraries.
pub struct CrfNER {
    /// Feature weights learned during training (or loaded from file)
    weights: HashMap<String, f64>,
    /// Entity type gazetteer lists
    gazetteers: HashMap<EntityType, Vec<String>>,
    /// Label set (BIO tagging)
    labels: Vec<String>,
    /// Feature templates
    templates: Vec<FeatureTemplate>,
}

/// Feature template for CRF
#[derive(Debug, Clone)]
pub enum FeatureTemplate {
    /// Current word
    Word,
    /// Word at offset
    WordAt(i32),
    /// Word shape (Xx, XX, x, 0)
    Shape,
    /// Shape at offset
    ShapeAt(i32),
    /// Prefix of length n
    Prefix(usize),
    /// Suffix of length n
    Suffix(usize),
    /// Is in gazetteer for entity type
    InGazetteer(EntityType),
    /// Previous label
    PrevLabel,
    /// Bigram: current + previous label
    LabelBigram,
    /// Word + Label combination
    WordLabel,
}

impl Default for CrfNER {
    fn default() -> Self {
        Self::new()
    }
}

mod algorithm;
// CRF algorithm: feature extraction, Viterbi decoding, weight loading (see algorithm.rs).
impl Model for CrfNER {
    fn extract_entities(&self, text: &str, _language: Option<&str>) -> Result<Vec<Entity>> {
        if text.trim().is_empty() {
            return Ok(vec![]);
        }

        let tokens = Self::tokenize(text);
        if tokens.is_empty() {
            return Ok(vec![]);
        }

        let labels = self.viterbi_decode(&tokens);
        let entities = self.labels_to_entities(text, &tokens, &labels);

        Ok(entities)
    }

    fn supported_types(&self) -> Vec<EntityType> {
        vec![
            EntityType::Person,
            EntityType::Organization,
            EntityType::Location,
            EntityType::Other("MISC".to_string()),
        ]
    }

    fn is_available(&self) -> bool {
        true // Always available (no external dependencies)
    }

    fn name(&self) -> &'static str {
        "crf"
    }

    fn description(&self) -> &'static str {
        "CRF-based NER (classical statistical method)"
    }

    fn capabilities(&self) -> crate::ModelCapabilities {
        crate::ModelCapabilities {
            batch_capable: true,
            optimal_batch_size: Some(32),
            streaming_capable: true,
            ..Default::default()
        }
    }
}

impl crate::NamedEntityCapable for CrfNER {}

impl crate::BatchCapable for CrfNER {
    fn optimal_batch_size(&self) -> Option<usize> {
        Some(32) // CRF is fast, can handle batches
    }
}

impl crate::StreamingCapable for CrfNER {
    fn recommended_chunk_size(&self) -> usize {
        4096 // Smaller chunks since CRF is token-based
    }
}

#[cfg(test)]
mod tests;