english_pos_tagger/

lib.rs

//! Part-of-speech tagging and lemmatization of English text.
//!
//! This crate assigns Penn-Treebank-style POS tags to tokens and adds base-form
//! lemmas for nouns (excluding proper nouns), verbs, and adjectives. It follows
//! a Brill-style transformation-based approach: a unigram pass proposes
//! candidate tags, ordered context rules rewrite them from neighboring tokens,
//! and a final pass lemmatizes and applies proper-noun and year fixups.
//!
//! Every token gets a normal form, and lemmas are always produced for the
//! parts of speech that carry them.
//!
//! # Example
//!
//! ```
//! use english_pos_tagger::Tagger;
//!
//! let tagger = Tagger::new();
//! let tokens = tagger.tag_sentence("A bear just crossed the road.");
//! assert_eq!(tokens[1].value, "bear");
//! assert_eq!(tokens[1].pos.as_deref(), Some("NN"));
//! assert_eq!(tokens[1].lemma.as_deref(), Some("bear"));
//! ```
//!
//! # Lexicon state
//!
//! [`Tagger::update_lexicon`] adds or replaces word entries on that tagger
//! instance. State stays local to the instance, not shared across taggers.

#![forbid(unsafe_code)]
#![warn(missing_docs)]

mod data;
mod lemmatize;
mod normalize;
mod rules;
mod token;
mod tokenizer;
mod unigram;

use std::collections::HashMap;

pub use token::Token;

// Contraction surface forms whose lemma is fixed regardless of POS.
const LEMMA_EXCEPTIONS: &[(&str, &str)] = &[
    ("ai", "be"),
    ("ca", "can"),
    ("sha", "shall"),
    ("'ll", "will"),
    ("wo", "will"),
    ("'ve", "have"),
    ("'m", "am"),
    ("'re", "be"),
    ("n't", "not"),
    ("'d", "would"),
];

// Year and mid-decade forms forced to CD.
const YEARS: &[&str] = &[
    "1990s",
    "1980s",
    "1970s",
    "1960s",
    "1950s",
    "1940s",
    "1930s",
    "1910s",
    "mid-1990s",
    "mid-1980s",
    "mid-1970s",
    "mid-1960s",
    "mid-1950s",
    "mid-1940s",
    "mid-1930s",
    "mid-1920s",
    "mid-1910s",
];

fn lemma_exception(word: &str) -> Option<&'static str> {
    LEMMA_EXCEPTIONS
        .iter()
        .find(|(k, _)| *k == word)
        .map(|(_, v)| *v)
}

fn is_year(word: &str) -> bool {
    YEARS.contains(&word)
}

/// A part-of-speech tagger with its own lexicon.
pub struct Tagger {
    lexicon: HashMap<String, Vec<String>>,
}

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

impl Tagger {
    /// Creates a tagger seeded with the built-in lexicon.
    pub fn new() -> Self {
        Tagger {
            lexicon: data::LEXICON.clone(),
        }
    }

    /// Adds or replaces lexicon entries. Each key normalizes before storage.
    /// Each value is a candidate POS list with the most frequent tag first.
    pub fn update_lexicon<I, K, V>(&mut self, entries: I)
    where
        I: IntoIterator<Item = (K, V)>,
        K: AsRef<str>,
        V: IntoIterator,
        V::Item: Into<String>,
    {
        for (key, value) in entries {
            let normalized = normalize::normalize(key.as_ref());
            let poses: Vec<String> = value.into_iter().map(Into::into).collect();
            self.lexicon.insert(normalized, poses);
        }
    }

    /// Tags `tokens` and returns them. Sets `normal`, `pos`, and, where
    /// applicable, `lemma` on each token.
    pub fn tag(&self, mut tokens: Vec<Token>) -> Vec<Token> {
        let mut poses: Vec<Vec<String>> = Vec::with_capacity(tokens.len());
        for token in tokens.iter_mut() {
            token.normal = Some(normalize::normalize(&token.value));
            poses.push(unigram::unigram_pos_tagger(token, &self.lexicon));
        }
        rules::apply_all_context_rules(&mut tokens, &poses);
        self.lemmatize(&mut tokens);
        tokens
    }

    /// Alias of [`Tagger::tag`].
    pub fn tag_tokens(&self, tokens: Vec<Token>) -> Vec<Token> {
        self.tag(tokens)
    }

    /// Classifies each raw string as number, punctuation, or word, then tags.
    pub fn tag_raw_tokens<S: AsRef<str>>(&self, raw_tokens: &[S]) -> Vec<Token> {
        let tokens: Vec<Token> = raw_tokens
            .iter()
            .map(|t| {
                let t = t.as_ref();
                let tag = if rgx::is_number(t) {
                    "number"
                } else if rgx::is_punctuation(t) {
                    "punctuation"
                } else {
                    "word"
                };
                Token::new(t, tag)
            })
            .collect();
        self.tag(tokens)
    }

    /// Tokenizes `sentence` with the default configuration, then tags it.
    pub fn tag_sentence(&self, sentence: &str) -> Vec<Token> {
        self.tag(tokenizer::tokenize(sentence))
    }

    // Lemmatizes tokens and applies the capitalized-NNP and year fixups.
    fn lemmatize(&self, tokens: &mut [Token]) {
        for token in tokens.iter_mut() {
            let w = token.normal.clone().unwrap_or_default();
            let v0 = token.value.chars().next();
            let cap = matches!(v0, Some(c) if c.is_ascii_uppercase());

            if is_year(&w) {
                token.pos = Some("CD".to_string());
            }

            if let Some(lemma) = lemma_exception(&w) {
                token.lemma = Some(lemma.to_string());
                continue;
            }

            let pos = match &token.pos {
                Some(p) => p.clone(),
                None => continue,
            };
            let first = pos.chars().next();
            match first {
                Some('J') => {
                    if cap {
                        token.lemma = Some(w.clone());
                        token.pos = Some("NNP".to_string());
                    } else if pos.len() > 2 {
                        token.lemma = Some(lemmatize::adjective(&w));
                    } else {
                        token.lemma = Some(w.clone());
                    }
                }
                Some('V') => {
                    if pos.len() > 2 {
                        let lemma = if w == "'s" {
                            "be".to_string()
                        } else {
                            lemmatize::verb(&w)
                        };
                        token.lemma = Some(lemma);
                    } else {
                        token.lemma = Some(w.clone());
                    }
                }
                Some('N') => {
                    if cap {
                        token.lemma = Some(w.clone());
                        token.pos = Some("NNP".to_string());
                    } else if pos != "NNP" && pos.len() > 2 {
                        token.lemma = Some(lemmatize::noun(&w));
                    } else {
                        token.lemma = Some(w.clone());
                    }
                }
                Some('M') => {
                    token.lemma = Some(lemmatize::verb(&w));
                }
                _ => {}
            }
        }
    }
}

// Raw-token classification regexes. Anchoring is deliberate.
mod rgx {
    use once_cell::sync::Lazy;
    use regex::Regex;

    static NUMBER: Lazy<Regex> =
        Lazy::new(|| Regex::new(r"(?-u:^(?:\d+/\d+|\d(?:[.,\-/]?\d)*(?:\.\d+)?)$)").unwrap());
    static PUNCTUATION: Lazy<Regex> = Lazy::new(|| {
        Regex::new("^[\u{2019}'\u{2018}\u{2019}`\u{201C}\u{201D}\"\\[\\](){}\u{2026},.!;?\\-:]+$")
            .unwrap()
    });

    pub fn is_number(s: &str) -> bool {
        NUMBER.is_match(s)
    }

    pub fn is_punctuation(s: &str) -> bool {
        PUNCTUATION.is_match(s)
    }
}