pos-tagger 0.2.0

//! Part-of-speech tagger.
//!
//! Tagging runs in two phases. First each word gets an initial tag from the
//! lexicon, falling back to a lowercase lookup and then to `NN`. Second the eight
//! contextual rules rewrite tags left to right.

use crate::lexicon;
use crate::rules;
use std::collections::HashMap;

/// A Brill-style English part-of-speech tagger.
///
/// Each tagger owns its lexicon, so [`Tagger::extend_lexicon`] affects only that
/// instance. Construct one with [`Tagger::new`].
///
/// ```
/// use pos_tagger::Tagger;
/// let tagger = Tagger::new();
/// let tagged = tagger.tag(&["the", "dog"]);
/// assert_eq!(tagged, vec![
///     ("the".to_string(), "DT".to_string()),
///     ("dog".to_string(), "NN".to_string()),
/// ]);
/// ```
#[derive(Debug, Clone)]
pub struct Tagger {
    lexicon: HashMap<String, Vec<String>>,
}

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

impl Tagger {
    /// Create a tagger with the built-in English lexicon.
    pub fn new() -> Self {
        Tagger {
            lexicon: lexicon::load(),
        }
    }

    /// Tag each word in `words`.
    ///
    /// Returns a vector the same length as the input. Each entry pairs the
    /// original word with its tag. The input is never modified.
    ///
    /// `words` is any slice whose elements borrow as `str`, so `&[&str]`,
    /// `&[String]`, and the lexer's `Vec<String>` output all pass without an
    /// intermediate `Vec<&str>`. Write `tagger.tag(&lexer.lex(text))` directly.
    ///
    /// Each result pair owns its word as a fresh `String`. The pair shape is the
    /// public contract, so the word is copied out of the borrowed input. Callers
    /// that only need tags can read the second field and drop the first.
    ///
    /// An empty token tags as `NN`. It is a lexicon key with an empty tag list,
    /// and the empty list falls back to `NN` instead of panicking.
    pub fn tag<S: AsRef<str>>(&self, words: &[S]) -> Vec<(String, String)> {
        let mut sentence: Vec<(String, String)> = words
            .iter()
            .map(|word| {
                let word = word.as_ref();
                (word.to_string(), self.initial_tag(word))
            })
            .collect();

        for index in 0..sentence.len() {
            rules::apply_all(&mut sentence, index);
        }
        sentence
    }

    /// Compute the initial tag for one word before the rules run.
    ///
    /// Look up the exact word first. On a miss, look up the lowercase form. On a
    /// second miss, default to `NN`. When a lookup hits, use the first tag in the
    /// candidate list.
    fn initial_tag(&self, word: &str) -> String {
        if let Some(tags) = self.lexicon.get(word) {
            return first_tag(tags);
        }
        if let Some(tags) = self.lexicon.get(&word.to_lowercase()) {
            return first_tag(tags);
        }
        "NN".to_string()
    }

    /// Report whether `word` is in the lexicon.
    ///
    /// Returns true when the exact word is a key or when its lowercase form is a
    /// key. The empty string is a key with an empty tag list, so
    /// `word_in_lexicon("")` returns true.
    ///
    /// The lexicon is a `HashMap`, so object member names such as `constructor`,
    /// `valueOf`, `hasOwnProperty`, `toString`, and `__proto__` return false
    /// unless they are real keys. None of them are. A lookup over a plain object
    /// would inherit these names and report true. This crate keeps the `HashMap`
    /// result on purpose.
    pub fn word_in_lexicon(&self, word: &str) -> bool {
        if self.lexicon.contains_key(word) {
            return true;
        }
        self.lexicon.contains_key(&word.to_lowercase())
    }

    /// Add entries from `entries`, keeping existing keys untouched.
    ///
    /// Only keys absent from the current lexicon are inserted. Existing entries
    /// are never overwritten.
    ///
    /// `entries` is any iterable of `(word, tags)` pairs. A `HashMap`, a
    /// `BTreeMap`, a `Vec`, an array, or a `map` adaptor all work, so callers do
    /// not build a `HashMap` first.
    pub fn extend_lexicon<I>(&mut self, entries: I)
    where
        I: IntoIterator<Item = (String, Vec<String>)>,
    {
        for (word, tags) in entries {
            self.lexicon.entry(word).or_insert(tags);
        }
    }
}

/// First tag of a candidate list, or `NN` when the list is empty.
///
/// The lexicon has one empty-tag entry, the key for the empty string. The lexer
/// never emits an empty token, so this path is unreachable through the normal
/// pipeline. Returning `NN` keeps a direct `tag(&[""])` call total instead of
/// panicking.
fn first_tag(tags: &[String]) -> String {
    tags.first().cloned().unwrap_or_else(|| "NN".to_string())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn unknown_word_defaults_to_nn() {
        let t = Tagger::new();
        assert_eq!(t.tag(&["Zzqxxw"]), vec![("Zzqxxw".into(), "NN".into())]);
    }

    #[test]
    fn lowercase_fallback() {
        let t = Tagger::new();
        assert_eq!(t.tag(&["THE"]), vec![("THE".into(), "DT".into())]);
    }
}