token-align 0.1.0

#![no_std]
#![forbid(unsafe_code)]
#![warn(missing_docs)]
//! Maps overlaps between token sequences and between character positions.
//!
//! The default comparison form removes common subword markers and applies
//! pinned Unicode compatibility decomposition. Use [`Exact`] or implement
//! [`Normalizer`] when the input needs different comparison rules.

extern crate alloc;

use alloc::string::String;
use alloc::vec;
use alloc::vec::Vec;

mod diff;
mod unicode_tables;

#[cfg(test)]
mod unit_tests;

/// One list of counterpart indices for every input token.
pub type Alignment = Vec<Vec<usize>>;

/// One list of counterpart indices for every input character.
pub type CharMap = Vec<Vec<usize>>;

/// Produces the comparison form of one token.
///
/// `out` is empty when this method starts. Implementations append the complete
/// comparison form and may leave it empty when the token has no counterpart.
/// A caller can also use a closure with the same arguments.
///
/// Byte-level token schemes often encode a non-ASCII character as several
/// visible characters. A normalizer for such a scheme can decode the token
/// before appending its comparison form.
///
/// # Examples
///
/// This normalizer decodes one visible byte-level spelling before alignment.
///
/// ```
/// use token_align::{get_alignments_with, Normalizer};
///
/// struct DecodeByteToken;
///
/// impl Normalizer for DecodeByteToken {
///     fn normalize(&self, token: &str, out: &mut String) {
///         if token == "é" {
///             out.push('é');
///         } else {
///             out.push_str(token);
///         }
///     }
/// }
///
/// let encoded = ["caf", "é"];
/// let text = ["café"];
/// assert_eq!(
///     get_alignments_with(&encoded, &text, &DecodeByteToken),
///     (vec![vec![0], vec![0]], vec![vec![0, 1]])
/// );
/// ```
pub trait Normalizer {
    /// Appends the comparison form of `token` to `out`.
    fn normalize(&self, token: &str, out: &mut String);
}

impl<F> Normalizer for F
where
    F: Fn(&str, &mut String),
{
    fn normalize(&self, token: &str, out: &mut String) {
        self(token, out);
    }
}

/// The default marker removal and Unicode comparison form.
///
/// The fixed control tokens `[CLS]`, `[SEP]`, `[PAD]`, `[UNK]`, `[MASK]`,
/// `<s>`, `</s>`, `<pad>`, `<unk>`, `<mask>`, `<cls>`, `<sep>`,
/// `<|endoftext|>`, and `<|startoftext|>` produce an empty comparison form.
///
/// One leading `##` and one trailing `</w>` are removed. Every `▁`, `Ġ`, and
/// `Ċ` is dropped. `Ġ` and `Ċ` are also Maltese capital letters, so literal
/// Maltese text containing either character loses it under this comparison
/// form. Use a custom [`Normalizer`] when those letters carry text.
///
/// Each remaining character receives recursive compatibility decomposition
/// from the Unicode Character Database 17.0.0 data. The decomposition is
/// canonically ordered within that character, then mapped with the simple
/// lowercase field. Final sigma `ς` maps to sigma `σ`. Characters are folded
/// independently, so marks are not reordered across input characters.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct Fold;

impl Normalizer for Fold {
    fn normalize(&self, token: &str, out: &mut String) {
        if is_special(token) {
            return;
        }

        let token = token.strip_prefix("##").unwrap_or(token);
        let token = token.strip_suffix("</w>").unwrap_or(token);
        for c in token.chars() {
            if !matches!(c, '\u{2581}' | '\u{0120}' | '\u{010A}') {
                unicode_tables::push_folded(c, out);
            }
        }
    }
}

/// A comparison form that copies each token unchanged.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct Exact;

impl Normalizer for Exact {
    fn normalize(&self, token: &str, out: &mut String) {
        out.push_str(token);
    }
}

/// Aligns two token sequences with the default [`Fold`] comparison form.
///
/// The first result maps `a` indices to `b` indices. The second result maps
/// `b` indices to `a` indices.
///
/// # Examples
///
/// ```
/// use token_align::get_alignments;
///
/// let a = ["New York"];
/// let b = ["New", "York"];
/// assert_eq!(
///     get_alignments(&a, &b),
///     (vec![vec![0, 1]], vec![vec![0], vec![0]])
/// );
/// ```
pub fn get_alignments<A, B>(a: &[A], b: &[B]) -> (Alignment, Alignment)
where
    A: AsRef<str>,
    B: AsRef<str>,
{
    get_alignments_with(a, b, &Fold)
}

/// Aligns two token sequences with a caller-provided comparison form.
pub fn get_alignments_with<A, B, N>(a: &[A], b: &[B], norm: &N) -> (Alignment, Alignment)
where
    A: AsRef<str>,
    B: AsRef<str>,
    N: Normalizer + ?Sized,
{
    let left = Stream::from_tokens(a, norm);
    let right = Stream::from_tokens(b, norm);
    align_streams(&left, &right, a.len(), b.len())
}

/// Maps Unicode scalar positions in `a` to positions in `b` and back.
///
/// # Examples
///
/// ```
/// use token_align::get_charmap;
///
/// assert_eq!(get_charmap("bar", "Bar"), (vec![vec![0], vec![1], vec![2]], vec![vec![0], vec![1], vec![2]]));
/// ```
pub fn get_charmap(a: &str, b: &str) -> (CharMap, CharMap) {
    get_charmap_with(a, b, &Fold)
}

/// Maps Unicode scalar positions with a caller-provided comparison form.
pub fn get_charmap_with<N>(a: &str, b: &str, norm: &N) -> (CharMap, CharMap)
where
    N: Normalizer + ?Sized,
{
    let left = Stream::from_chars(a, norm);
    let right = Stream::from_chars(b, norm);
    align_streams(&left, &right, a.chars().count(), b.chars().count())
}

fn is_special(token: &str) -> bool {
    matches!(
        token,
        "[CLS]"
            | "[SEP]"
            | "[PAD]"
            | "[UNK]"
            | "[MASK]"
            | "<s>"
            | "</s>"
            | "<pad>"
            | "<unk>"
            | "<mask>"
            | "<cls>"
            | "<sep>"
            | "<|endoftext|>"
            | "<|startoftext|>"
    )
}

struct Stream {
    chars: Vec<char>,
    owners: Vec<usize>,
    units: Vec<Unit>,
}

impl Stream {
    fn from_tokens<T, N>(tokens: &[T], norm: &N) -> Self
    where
        T: AsRef<str>,
        N: Normalizer + ?Sized,
    {
        let estimated = tokens.iter().fold(0usize, |sum, token| {
            sum.saturating_add(token.as_ref().len())
        });
        let mut stream = Self {
            chars: Vec::with_capacity(estimated),
            owners: Vec::with_capacity(estimated),
            units: Vec::new(),
        };
        let mut scratch = String::new();
        for (index, token) in tokens.iter().enumerate() {
            scratch.clear();
            norm.normalize(token.as_ref(), &mut scratch);
            stream.extend(&scratch, index);
        }
        stream.group_units();
        stream
    }

    fn from_chars<N>(text: &str, norm: &N) -> Self
    where
        N: Normalizer + ?Sized,
    {
        let count = text.chars().count();
        let mut stream = Self {
            chars: Vec::with_capacity(count),
            owners: Vec::with_capacity(count),
            units: Vec::new(),
        };
        let mut scratch = String::new();
        for (index, (start, c)) in text.char_indices().enumerate() {
            let end = start + c.len_utf8();
            scratch.clear();
            norm.normalize(&text[start..end], &mut scratch);
            stream.extend(&scratch, index);
        }
        stream.group_units();
        stream
    }

    fn extend(&mut self, text: &str, owner: usize) {
        for c in text.chars() {
            self.chars.push(c);
            self.owners.push(owner);
        }
    }

    fn group_units(&mut self) {
        for (index, &c) in self.chars.iter().enumerate() {
            if self.units.is_empty() || !unicode_tables::is_extender(c) {
                self.units.push(Unit {
                    start: index,
                    end: index + 1,
                });
            } else if let Some(unit) = self.units.last_mut() {
                unit.end = index + 1;
            }
        }
    }

    fn keys(&self) -> impl ExactSizeIterator<Item = char> + '_ {
        self.units.iter().map(|unit| self.chars[unit.start])
    }
}

#[derive(Clone, Copy)]
struct Unit {
    start: usize,
    end: usize,
}

fn align_streams(
    left: &Stream,
    right: &Stream,
    left_len: usize,
    right_len: usize,
) -> (Alignment, Alignment) {
    let left_keys: Vec<char> = left.keys().collect();
    let right_keys: Vec<char> = right.keys().collect();
    let mut unit_pairs = Vec::with_capacity(left_keys.len().min(right_keys.len()));
    diff::matched_pairs(&left_keys, &right_keys, &mut unit_pairs);

    let mut left_to_right = vec![Vec::new(); left_len];
    let mut right_to_left = vec![Vec::new(); right_len];
    for (left_unit, right_unit) in unit_pairs {
        project_unit(
            left,
            left_unit,
            right,
            right_unit,
            &mut left_to_right,
            &mut right_to_left,
        );
    }
    (left_to_right, right_to_left)
}

fn project_unit(
    left: &Stream,
    left_index: usize,
    right: &Stream,
    right_index: usize,
    left_to_right: &mut Alignment,
    right_to_left: &mut Alignment,
) {
    let left_unit = left.units[left_index];
    let right_unit = right.units[right_index];
    project_pair(
        left,
        left_unit.start,
        right,
        right_unit.start,
        left_to_right,
        right_to_left,
    );

    let mut right_cursor = right_unit.start + 1;
    for left_char in left_unit.start + 1..left_unit.end {
        while right_cursor < right_unit.end && right.chars[right_cursor] != left.chars[left_char] {
            right_cursor += 1;
        }
        if right_cursor < right_unit.end {
            project_pair(
                left,
                left_char,
                right,
                right_cursor,
                left_to_right,
                right_to_left,
            );
            right_cursor += 1;
        }
    }
}

fn project_pair(
    left: &Stream,
    left_char: usize,
    right: &Stream,
    right_char: usize,
    left_to_right: &mut Alignment,
    right_to_left: &mut Alignment,
) {
    let left_owner = left.owners[left_char];
    let right_owner = right.owners[right_char];
    if left_owner >= left_to_right.len() || right_owner >= right_to_left.len() {
        return;
    }
    push_new(&mut left_to_right[left_owner], right_owner);
    push_new(&mut right_to_left[right_owner], left_owner);
}

fn push_new(indices: &mut Vec<usize>, index: usize) {
    if indices.last().copied() != Some(index) {
        indices.push(index);
    }
}