use std::fmt;
use super::adapt::{adapt, adapt_with_df};
use super::bag::{BagKind, BagOptions, bag};
use super::error::TextError;
use super::limits::TextLimits;
use super::split::Split;
use super::standardize::Standardize;
use super::transform::{TransformOptions, transform};
use super::vocab::{
RESERVED_COUNT, Vocabulary, check_max_tokens, check_vocabulary_within_max_tokens,
};
use crate::Tensor;
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct TextVectorizationConfig {
pub max_tokens: Option<usize>,
pub standardize: Standardize,
pub split: Split,
pub ngrams: Option<usize>,
pub output_sequence_length: Option<usize>,
pub pad_to_max_tokens: bool,
pub limits: TextLimits,
}
impl TextVectorizationConfig {
pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
self.max_tokens = Some(max_tokens);
self
}
pub fn with_standardize(mut self, standardize: Standardize) -> Self {
self.standardize = standardize;
self
}
pub fn with_split(mut self, split: Split) -> Self {
self.split = split;
self
}
pub fn with_ngrams(mut self, n: usize) -> Self {
self.ngrams = Some(n);
self
}
pub fn with_output_sequence_length(mut self, len: usize) -> Self {
self.output_sequence_length = Some(len);
self
}
pub fn with_pad_to_max_tokens(mut self, pad: bool) -> Self {
self.pad_to_max_tokens = pad;
self
}
pub fn with_limits(mut self, limits: TextLimits) -> Self {
self.limits = limits;
self
}
}
#[derive(Clone)]
pub struct TextVectorization {
config: TextVectorizationConfig,
vocabulary: Vocabulary,
idf: Option<Vec<f32>>,
}
fn check_pad(config: &TextVectorizationConfig) -> Result<(), TextError> {
if config.pad_to_max_tokens && config.max_tokens.is_none() {
return Err(TextError::PadToMaxTokensWithoutMaxTokens);
}
Ok(())
}
fn assemble_idf(vocab_weights: &[f32]) -> Vec<f32> {
let oov = if vocab_weights.is_empty() {
0.0f32
} else {
let sum: f64 = vocab_weights.iter().map(|&w| f64::from(w)).sum();
(sum / vocab_weights.len() as f64) as f32
};
let mut idf = Vec::with_capacity(vocab_weights.len() + RESERVED_COUNT);
idf.push(0.0f32);
idf.push(oov);
idf.extend_from_slice(vocab_weights);
idf
}
fn check_config_numbers(config: &TextVectorizationConfig) -> Result<(), TextError> {
if let Some(n) = config.ngrams {
config.limits.check_ngrams(n)?;
}
if let Some(l) = config.output_sequence_length {
config.limits.check_output_sequence_length(l)?;
}
Ok(())
}
impl TextVectorization {
pub fn from_vocabulary<S: AsRef<str>>(
config: TextVectorizationConfig,
vocabulary: &[S],
) -> Result<Self, TextError> {
check_pad(&config)?;
check_max_tokens(config.max_tokens)?;
check_vocabulary_within_max_tokens(config.max_tokens, vocabulary.len())?;
check_config_numbers(&config)?;
let vocabulary = Vocabulary::from_tokens(&config.limits, vocabulary)?;
Ok(Self {
config,
vocabulary,
idf: None,
})
}
pub fn from_vocabulary_with_idf<S: AsRef<str>>(
config: TextVectorizationConfig,
vocabulary: &[S],
idf_weights: &[f32],
) -> Result<Self, TextError> {
check_pad(&config)?;
check_max_tokens(config.max_tokens)?;
check_vocabulary_within_max_tokens(config.max_tokens, vocabulary.len())?;
check_config_numbers(&config)?;
if idf_weights.len() != vocabulary.len() {
return Err(TextError::IdfWeightsLengthMismatch {
len: idf_weights.len(),
expected: vocabulary.len(),
});
}
if let Some(index) = idf_weights
.iter()
.position(|w| !w.is_finite() || w.is_sign_negative() && *w != 0.0)
{
return Err(TextError::InvalidIdfWeight { index });
}
let vocabulary = Vocabulary::from_tokens(&config.limits, vocabulary)?;
let idf = Some(assemble_idf(idf_weights));
Ok(Self {
config,
vocabulary,
idf,
})
}
pub fn adapt<S: AsRef<str>>(
config: TextVectorizationConfig,
corpus: &[S],
) -> Result<Self, TextError> {
check_pad(&config)?;
check_max_tokens(config.max_tokens)?;
check_config_numbers(&config)?;
let vocabulary = adapt(
corpus,
config.standardize,
config.split,
config.ngrams,
config.max_tokens,
&config.limits,
)?;
Ok(Self {
config,
vocabulary,
idf: None,
})
}
pub fn adapt_tf_idf<S: AsRef<str>>(
config: TextVectorizationConfig,
corpus: &[S],
) -> Result<Self, TextError> {
check_pad(&config)?;
check_max_tokens(config.max_tokens)?;
check_config_numbers(&config)?;
let (vocabulary, df) = adapt_with_df(
corpus,
config.standardize,
config.split,
config.ngrams,
config.max_tokens,
&config.limits,
)?;
let n = corpus.len() as f64;
let weights: Vec<f32> = df
.iter()
.map(|&d| (1.0f64 + n / (1.0f64 + d as f64)).ln() as f32)
.collect();
let idf = Some(assemble_idf(&weights));
Ok(Self {
config,
vocabulary,
idf,
})
}
pub fn transform<S: AsRef<str>>(&self, inputs: &[S]) -> Result<Tensor<i32>, TextError> {
let opts = TransformOptions {
standardize: self.config.standardize,
split: self.config.split,
ngrams: self.config.ngrams,
output_sequence_length: self.config.output_sequence_length,
};
transform(&self.vocabulary, &opts, &self.config.limits, inputs)
}
fn bag_width(&self) -> usize {
match (self.config.pad_to_max_tokens, self.config.max_tokens) {
(true, Some(m)) => m,
_ => self.vocabulary.vocabulary_size(),
}
}
fn bag_options(&self) -> BagOptions {
BagOptions {
standardize: self.config.standardize,
split: self.config.split,
ngrams: self.config.ngrams,
width: self.bag_width(),
}
}
pub fn transform_multi_hot<S: AsRef<str>>(
&self,
inputs: &[S],
) -> Result<Tensor<f32>, TextError> {
bag(
&self.vocabulary,
&self.bag_options(),
&self.config.limits,
inputs,
BagKind::Presence,
)
}
pub fn transform_count<S: AsRef<str>>(&self, inputs: &[S]) -> Result<Tensor<f32>, TextError> {
bag(
&self.vocabulary,
&self.bag_options(),
&self.config.limits,
inputs,
BagKind::Count,
)
}
pub fn transform_tf_idf<S: AsRef<str>>(&self, inputs: &[S]) -> Result<Tensor<f32>, TextError> {
let Some(idf) = self.idf.as_deref() else {
return Err(TextError::MissingIdfWeights);
};
bag(
&self.vocabulary,
&self.bag_options(),
&self.config.limits,
inputs,
BagKind::Weighted(idf),
)
}
pub fn idf_weights(&self) -> Option<&[f32]> {
self.idf.as_deref()
}
pub fn vocabulary(&self) -> &[String] {
self.vocabulary.tokens()
}
pub fn vocabulary_size(&self) -> usize {
self.vocabulary.vocabulary_size()
}
pub fn config(&self) -> &TextVectorizationConfig {
&self.config
}
}
impl fmt::Debug for TextVectorization {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TextVectorization")
.field("config", &self.config)
.field("vocabulary_size", &self.vocabulary.vocabulary_size())
.field("idf_weights_len", &self.idf.as_ref().map(Vec::len))
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_config_matches_recorded_defaults() {
let c = TextVectorizationConfig::default();
assert_eq!(c.max_tokens, None);
assert_eq!(c.standardize, Standardize::LowerAndStripPunctuation);
assert_eq!(c.split, Split::Whitespace);
assert_eq!(c.ngrams, None);
assert_eq!(c.output_sequence_length, None);
assert_eq!(c.limits, TextLimits::default());
}
#[test]
fn builders_set_some_and_fields() {
let c = TextVectorizationConfig::default()
.with_max_tokens(10)
.with_standardize(Standardize::None)
.with_split(Split::Character)
.with_ngrams(2)
.with_output_sequence_length(5);
assert_eq!(c.max_tokens, Some(10));
assert_eq!(c.standardize, Standardize::None);
assert_eq!(c.split, Split::Character);
assert_eq!(c.ngrams, Some(2));
assert_eq!(c.output_sequence_length, Some(5));
}
#[test]
fn from_vocabulary_max_tokens_is_fail_closed() {
let cfg = |m| TextVectorizationConfig::default().with_max_tokens(m);
let v = ["a", "b", "c"];
assert!(matches!(
TextVectorization::from_vocabulary(cfg(2), &v),
Err(TextError::InvalidMaxTokens { max_tokens: 2 })
));
assert!(matches!(
TextVectorization::from_vocabulary(cfg(4), &v),
Err(TextError::VocabularyTooLarge { len: 5, max: 4 })
));
assert_eq!(
TextVectorization::from_vocabulary(cfg(5), &v)
.unwrap()
.vocabulary_size(),
5
);
}
#[test]
fn construction_checks_numbers_before_vocabulary() {
let empty: [&str; 0] = [];
let bad_n = TextVectorizationConfig::default().with_ngrams(0);
assert!(matches!(
TextVectorization::from_vocabulary(bad_n.clone(), &empty),
Err(TextError::InvalidNgrams { .. })
));
assert!(matches!(
TextVectorization::adapt(bad_n, &empty),
Err(TextError::InvalidNgrams { .. })
));
let too_long = TextVectorizationConfig::default().with_output_sequence_length(usize::MAX);
assert!(matches!(
TextVectorization::from_vocabulary(too_long.clone(), &empty),
Err(TextError::OutputSequenceLengthTooLarge { .. })
));
assert!(matches!(
TextVectorization::adapt(too_long, &empty),
Err(TextError::OutputSequenceLengthTooLarge { .. })
));
}
#[test]
fn debug_hides_vocabulary_contents() {
let tv = TextVectorization::from_vocabulary(
TextVectorizationConfig::default(),
&["SECRET_MARKER_TOKEN"],
)
.unwrap();
let s = format!("{tv:?}");
assert!(!s.contains("SECRET_MARKER_TOKEN"));
assert!(s.contains("vocabulary_size"));
}
fn idf_of(n: usize, df: usize) -> f32 {
(1.0f64 + n as f64 / (1.0 + df as f64)).ln() as f32
}
fn oov_mean(ws: &[f32]) -> f32 {
if ws.is_empty() {
return 0.0;
}
(ws.iter().map(|&w| f64::from(w)).sum::<f64>() / ws.len() as f64) as f32
}
#[test]
fn pad_field_defaults_false_and_builder_sets_it() {
assert!(!TextVectorizationConfig::default().pad_to_max_tokens);
let c = TextVectorizationConfig::default().with_pad_to_max_tokens(true);
assert!(c.pad_to_max_tokens);
}
#[test]
fn pad_without_max_tokens_is_rejected_by_all_constructors() {
let c = TextVectorizationConfig::default().with_pad_to_max_tokens(true);
let v = ["a"];
assert!(matches!(
TextVectorization::from_vocabulary(c.clone(), &v),
Err(TextError::PadToMaxTokensWithoutMaxTokens)
));
assert!(matches!(
TextVectorization::adapt(c.clone(), &v),
Err(TextError::PadToMaxTokensWithoutMaxTokens)
));
assert!(matches!(
TextVectorization::from_vocabulary_with_idf(c.clone(), &v, &[1.0]),
Err(TextError::PadToMaxTokensWithoutMaxTokens)
));
assert!(matches!(
TextVectorization::adapt_tf_idf(c, &v),
Err(TextError::PadToMaxTokensWithoutMaxTokens)
));
}
#[test]
fn width_follows_vocabulary_or_max_tokens_when_padded() {
let cfg = TextVectorizationConfig::default().with_max_tokens(8);
let tv = TextVectorization::from_vocabulary(cfg.clone(), &["a", "b"]).unwrap();
assert_eq!(tv.transform_count(&["a"]).unwrap().shape(), &[1, 4]);
let padded = cfg.with_pad_to_max_tokens(true);
let tv = TextVectorization::from_vocabulary(padded, &["a", "b"]).unwrap();
let out = tv.transform_count(&["a b b"]).unwrap();
assert_eq!(out.shape(), &[1, 8]);
assert_eq!(
out.host_slice().to_vec(),
[0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 0.0, 0.0]
);
assert_eq!(
tv.transform_multi_hot::<&str>(&[]).unwrap().shape(),
&[0, 8]
);
}
#[test]
fn output_sequence_length_is_ignored_by_bag_methods() {
let a = TextVectorizationConfig::default();
let b = a.clone().with_output_sequence_length(1);
let v = ["x", "y"];
let ta = TextVectorization::from_vocabulary(a, &v).unwrap();
let tb = TextVectorization::from_vocabulary(b, &v).unwrap();
let ins = ["x y y z"];
assert_eq!(
ta.transform_count(&ins).unwrap().host_slice(),
tb.transform_count(&ins).unwrap().host_slice()
);
}
#[test]
fn tf_idf_without_weights_is_missing_even_for_bad_input() {
let tv =
TextVectorization::from_vocabulary(TextVectorizationConfig::default(), &["a"]).unwrap();
assert!(tv.idf_weights().is_none());
assert!(matches!(
tv.transform_tf_idf(&["a"]),
Err(TextError::MissingIdfWeights)
));
let tv = TextVectorization::adapt(TextVectorizationConfig::default(), &["a"]).unwrap();
assert!(matches!(
tv.transform_tf_idf::<&str>(&[]),
Err(TextError::MissingIdfWeights)
));
}
#[test]
fn idf_weights_are_validated_in_documented_order() {
let cfg = TextVectorizationConfig::default;
assert!(matches!(
TextVectorization::from_vocabulary_with_idf(cfg(), &["a", "b"], &[1.0]),
Err(TextError::IdfWeightsLengthMismatch {
len: 1,
expected: 2
})
));
for (i, bad) in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY, -1.0]
.into_iter()
.enumerate()
{
let r = TextVectorization::from_vocabulary_with_idf(cfg(), &["a", "b"], &[0.5, bad]);
assert!(
matches!(r, Err(TextError::InvalidIdfWeight { index: 1 })),
"{i}"
);
}
assert!(
TextVectorization::from_vocabulary_with_idf(cfg(), &["a", "b"], &[-0.0, 0.0]).is_ok()
);
assert!(matches!(
TextVectorization::from_vocabulary_with_idf(cfg(), &["a", "a"], &[f32::NAN]),
Err(TextError::IdfWeightsLengthMismatch { .. })
));
assert!(matches!(
TextVectorization::from_vocabulary_with_idf(cfg(), &["a", "a"], &[f32::NAN, 1.0]),
Err(TextError::InvalidIdfWeight { index: 0 })
));
let small = cfg().with_max_tokens(3);
assert!(matches!(
TextVectorization::from_vocabulary_with_idf(small, &["a", "b"], &[f32::NAN]),
Err(TextError::VocabularyTooLarge { .. })
));
let pad = cfg().with_pad_to_max_tokens(true);
assert!(matches!(
TextVectorization::from_vocabulary_with_idf(pad, &["a"], &[f32::NAN, 1.0]),
Err(TextError::PadToMaxTokensWithoutMaxTokens)
));
}
#[test]
fn explicit_idf_layout_and_oov_mean() {
let tv = TextVectorization::from_vocabulary_with_idf(
TextVectorizationConfig::default(),
&["a", "b"],
&[1.0, 2.0],
)
.unwrap();
assert_eq!(tv.idf_weights().unwrap(), [0.0, 1.5, 1.0, 2.0]);
let empty: [&str; 0] = [];
let tv = TextVectorization::from_vocabulary_with_idf(
TextVectorizationConfig::default(),
&empty,
&[],
)
.unwrap();
assert_eq!(tv.idf_weights().unwrap(), [0.0, 0.0]);
}
#[test]
fn adapt_tf_idf_matches_formula_bit_for_bit_and_round_trips() {
let corpus = ["a a b", "a c", "b b b d"];
let cfg = TextVectorizationConfig::default();
let tv = TextVectorization::adapt_tf_idf(cfg.clone(), &corpus).unwrap();
assert_eq!(tv.vocabulary(), ["", "[UNK]", "b", "a", "c", "d"]);
let dfs = [2usize, 2, 1, 1];
let ws: Vec<f32> = dfs.iter().map(|&d| idf_of(3, d)).collect();
let idf = tv.idf_weights().unwrap();
assert_eq!(idf.len(), tv.vocabulary_size());
assert_eq!(idf[0].to_bits(), 0.0f32.to_bits());
assert_eq!(idf[1].to_bits(), oov_mean(&ws).to_bits());
for (g, w) in idf[2..].iter().zip(&ws) {
assert_eq!(g.to_bits(), w.to_bits());
}
let out = tv.transform_tf_idf(&["b b a zzz"]).unwrap();
let o = out.host_slice();
assert_eq!(o[0], 0.0);
assert_eq!(o[1].to_bits(), (1.0f32 * oov_mean(&ws)).to_bits());
assert_eq!(o[2].to_bits(), (2.0f32 * ws[0]).to_bits());
assert_eq!(o[3].to_bits(), (1.0f32 * ws[1]).to_bits());
assert_eq!(o[4], 0.0);
let words: Vec<String> = tv.vocabulary()[2..].to_vec();
let back = TextVectorization::from_vocabulary_with_idf(cfg, &words, &idf[2..]).unwrap();
let bi = back.idf_weights().unwrap();
assert!(idf.iter().zip(bi).all(|(a, b)| a.to_bits() == b.to_bits()));
let again = back.transform_tf_idf(&["b b a zzz"]).unwrap();
assert!(
o.iter()
.zip(again.host_slice().iter())
.all(|(a, b)| a.to_bits() == b.to_bits())
);
}
#[test]
fn tf_idf_never_reads_idf_for_pad_columns() {
let cfg = TextVectorizationConfig::default()
.with_max_tokens(6)
.with_pad_to_max_tokens(true);
let tv = TextVectorization::from_vocabulary_with_idf(cfg, &["a"], &[2.0]).unwrap();
let out = tv.transform_tf_idf(&["a a q"]).unwrap();
assert_eq!(out.shape(), &[1, 6]);
assert_eq!(out.host_slice().to_vec(), [0.0, 2.0, 4.0, 0.0, 0.0, 0.0]);
}
#[test]
fn debug_shows_only_idf_count() {
let tv = TextVectorization::from_vocabulary_with_idf(
TextVectorizationConfig::default(),
&["SECRET_MARKER_TOKEN"],
&[0.123_456_79],
)
.unwrap();
let s = format!("{tv:?}");
assert!(!s.contains("SECRET_MARKER_TOKEN"));
assert!(!s.contains("0.1234"));
assert!(s.contains("idf_weights_len: Some(3)"));
}
}