pub mod embedding;
pub mod error;
mod token_index;
pub use embedding::Embedding;
pub use error::{
ArtifactTokenizerRead, ContentlessInputOverBudget, ContractMismatch, EmbeddingDimMismatch, Error,
InputTooLarge, OutputShape, SpecialTokenOverhead, TokenCount, TokenizerContractMismatch,
WindowOverBudget,
};
pub use windit::plan::WindowOptions;
pub use windit::TailPolicy;
pub use crate::embeddings::long_text::LongTextOptions;
use std::{ops::Range, path::Path, sync::OnceLock};
use crate::{ComputeUnits, DataType, Model, MultiArray};
use tokenizers::{
PostProcessor, Tokenizer, TruncationDirection, TruncationParams, TruncationStrategy,
};
use crate::embeddings::granite::{
embedding::{EMBEDDING_DIM, check_finite_output},
error::Result,
token_index::{IndexMeasure, LazyTable, MergeTable, TokenIndex},
};
use crate::embeddings::long_text;
pub const TOKENIZER_FILE_NAME: &str = "tokenizer.json";
mod names {
pub const INPUT_IDS: &str = "input_ids";
pub const ATTENTION_MASK: &str = "attention_mask";
pub const EMBEDDING: &str = "embedding";
}
mod contract {
pub const VOCAB_SIZE: usize = 180_000;
pub const MAX_TOKEN_ID: u32 = 179_999;
pub const CLS_TOKEN: &str = "<|startoftext|>";
pub const CLS_ID: u32 = 179_934;
pub const PAD_TOKEN: &str = "<|endoftext|>";
pub const PAD_ID: u32 = 179_935;
pub const EOS_TOKEN: &str = "<|return|>";
pub const EOS_ID: u32 = 179_938;
pub const SENTINEL_TEXT: &str = "hello world";
pub const SENTINEL_IDS: [u32; 4] = [CLS_ID, 24_313, 2_318, EOS_ID];
pub const TOKENIZER_SHA256_HEX: &str =
"4f2842d568e2724370aec203652a42ac783c7937f8347a1a2cc7506d71f1582f";
}
pub const MAX_TOKENS: usize = 512;
const _: () = assert!(long_text::DEFAULT_WINDOW == MAX_TOKENS);
pub const SPECIAL_TOKENS_PER_WINDOW: usize = 2;
pub const CONTENT_TOKENS_PER_WINDOW: usize = MAX_TOKENS - SPECIAL_TOKENS_PER_WINDOW;
pub const DEFAULT_COMPUTE: ComputeUnits = ComputeUnits::All;
#[cfg(feature = "serde")]
fn default_compute() -> ComputeUnits {
DEFAULT_COMPUTE
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TextEmbedderOptions {
#[cfg_attr(feature = "serde", serde(default = "default_compute"))]
compute: ComputeUnits,
}
impl Default for TextEmbedderOptions {
fn default() -> Self {
Self::new()
}
}
impl TextEmbedderOptions {
pub const fn new() -> Self {
Self {
compute: DEFAULT_COMPUTE,
}
}
#[inline]
pub const fn compute(&self) -> ComputeUnits {
self.compute
}
#[must_use]
#[inline]
pub const fn with_compute(mut self, compute: ComputeUnits) -> Self {
self.set_compute(compute);
self
}
#[inline]
pub const fn set_compute(&mut self, compute: ComputeUnits) -> &mut Self {
self.compute = compute;
self
}
}
#[derive(Clone, Debug)]
pub struct WindowEmbedding {
ordinal: usize,
byte_range: Range<usize>,
token_span: windit::plan::Span,
embedding: Embedding,
}
impl WindowEmbedding {
#[inline]
pub const fn ordinal(&self) -> usize {
self.ordinal
}
#[inline]
pub const fn byte_start(&self) -> usize {
self.byte_range.start
}
#[inline]
pub const fn byte_end(&self) -> usize {
self.byte_range.end
}
#[inline]
pub fn byte_range(&self) -> Range<usize> {
self.byte_range.clone()
}
#[inline]
pub const fn token_span(&self) -> windit::plan::Span {
self.token_span
}
#[inline]
pub const fn token_count(&self) -> usize {
self.token_span.len()
}
#[inline]
pub const fn embedding(&self) -> &Embedding {
&self.embedding
}
#[must_use]
#[inline]
pub fn into_embedding(self) -> Embedding {
self.embedding
}
}
#[derive(Debug)]
pub struct TextEmbedder {
model: Model,
tokenizer: Tokenizer,
pad_id: i32,
measure_tokenizer: OnceLock<Tokenizer>,
merge_table: OnceLock<Option<MergeTable>>,
}
impl TextEmbedder {
pub fn load(model_path: impl AsRef<Path>, options: TextEmbedderOptions) -> Result<Self> {
let model_path = model_path.as_ref();
let tokenizer_path = artifact_tokenizer_path(model_path);
let bytes = std::fs::read(&tokenizer_path).map_err(|source| {
Error::ArtifactTokenizerRead(ArtifactTokenizerRead::new(tokenizer_path.clone(), source))
})?;
let sha256_hex = sha256_hex(&bytes);
validate_tokenizer_identity(&TokenizerProvenance::Artifact(Artifact::new(
tokenizer_path,
sha256_hex,
)))?;
let tokenizer = Tokenizer::from_bytes(&bytes).map_err(Error::TokenizerLoad)?;
Self::from_parts(model_path, tokenizer, options)
}
pub fn from_file(model_path: impl AsRef<Path>) -> Result<Self> {
Self::load(model_path, TextEmbedderOptions::new())
}
pub fn from_files(
model_path: impl AsRef<Path>,
tokenizer_json_path: impl AsRef<Path>,
options: TextEmbedderOptions,
) -> Result<Self> {
let bytes =
std::fs::read(tokenizer_json_path.as_ref()).map_err(|e| Error::TokenizerLoad(e.into()))?;
Self::from_memory(model_path, &bytes, options)
}
pub fn from_memory(
model_path: impl AsRef<Path>,
tokenizer_json_bytes: &[u8],
options: TextEmbedderOptions,
) -> Result<Self> {
let sha256_hex = sha256_hex(tokenizer_json_bytes);
validate_tokenizer_identity(&TokenizerProvenance::Supplied(sha256_hex))?;
let tokenizer = Tokenizer::from_bytes(tokenizer_json_bytes).map_err(Error::TokenizerLoad)?;
Self::from_parts(model_path, tokenizer, options)
}
fn from_parts(
model_path: impl AsRef<Path>,
mut tokenizer: Tokenizer,
options: TextEmbedderOptions,
) -> Result<Self> {
configure_tokenizer(&mut tokenizer)?;
validate_tokenizer_contract(&tokenizer)?;
let pad_id = tokenizer
.token_to_id("<|endoftext|>")
.and_then(|id| i32::try_from(id).ok())
.unwrap_or(0);
let model = Model::load(model_path, options.compute())?;
let description = model.description();
let ids_expected = format!("[1, {MAX_TOKENS}] int32");
for name in [names::INPUT_IDS, names::ATTENTION_MASK] {
let input = description.input(name).ok_or_else(|| {
Error::ContractMismatch(ContractMismatch::new(
name,
ids_expected.clone(),
"missing".to_string(),
))
})?;
if input.shape() != [1, MAX_TOKENS] || input.data_type() != Some(DataType::I32) {
return Err(Error::ContractMismatch(ContractMismatch::new(
name,
ids_expected.clone(),
describe(input.shape(), input.data_type()),
)));
}
}
let output_expected = format!("[1, {EMBEDDING_DIM}] float32");
let output = description.output(names::EMBEDDING).ok_or_else(|| {
Error::ContractMismatch(ContractMismatch::new(
names::EMBEDDING,
output_expected.clone(),
"missing".to_string(),
))
})?;
if output.shape() != [1, EMBEDDING_DIM] || output.data_type() != Some(DataType::F32) {
return Err(Error::ContractMismatch(ContractMismatch::new(
names::EMBEDDING,
output_expected,
describe(output.shape(), output.data_type()),
)));
}
Ok(Self {
model,
tokenizer,
pad_id,
measure_tokenizer: OnceLock::new(),
merge_table: OnceLock::new(),
})
}
pub fn token_ids(&self, text: &str) -> Result<Vec<u32>> {
if text.is_empty() {
return Err(Error::EmptyText);
}
let encoding = self.tokenizer.encode(text, true).map_err(Error::Tokenize)?;
Ok(encoding.get_ids().to_vec())
}
pub fn embed(&self, text: &str) -> Result<Embedding> {
let ids = self.token_ids(text)?;
self.embed_tokenized(&ids)
}
fn embed_tokenized(&self, ids: &[u32]) -> Result<Embedding> {
let (input_ids, attention_mask) = build_window(ids, self.pad_id)?;
let ids_tensor = MultiArray::from_slice(&[1, MAX_TOKENS], &input_ids)?;
let mask_tensor = MultiArray::from_slice(&[1, MAX_TOKENS], &attention_mask)?;
let mut outputs = self.model.predict_with(&[
(names::INPUT_IDS, &ids_tensor),
(names::ATTENTION_MASK, &mask_tensor),
])?;
let embeds = outputs
.take(names::EMBEDDING)
.ok_or_else(|| crate::PredictionError::MissingOutput(names::EMBEDDING.to_string()))?;
if embeds.shape() != [1, EMBEDDING_DIM] {
return Err(Error::OutputShape(OutputShape::new(
embeds.shape().to_vec(),
vec![1, EMBEDDING_DIM],
)));
}
let mut row = [0.0f32; EMBEDDING_DIM];
embeds.copy_into::<f32>(&mut row)?;
check_finite_output(&row)?;
Embedding::from_slice_normalizing(&row)
}
fn measuring_tokenizer(&self) -> Result<&Tokenizer> {
if let Some(t) = self.measure_tokenizer.get() {
return Ok(t);
}
let mut t = self.tokenizer.clone();
t.with_truncation(None).map_err(Error::TokenizerConfig)?;
let _ = self.measure_tokenizer.set(t);
Ok(
self
.measure_tokenizer
.get()
.expect("measure_tokenizer was set just above, on this thread or another"),
)
}
pub fn embed_long(&self, text: &str) -> Result<Embedding> {
self.embed_long_with(text, &LongTextOptions::new())
}
pub fn embed_long_with(&self, text: &str, opts: &LongTextOptions) -> Result<Embedding> {
let mut windows = self.embed_windows_with(text, opts)?;
if windows.len() == 1 {
return Ok(
windows
.pop()
.expect("a one-element vector has a last element")
.into_embedding(),
);
}
let windowed: Vec<_> = windows
.into_iter()
.map(|w| windit::windowed::Windowed::new(w.embedding, w.token_span))
.collect();
Ok(windit::aggregate::aggregate(
&windit::aggregate::CoverageWeightedMean,
&windowed,
)?)
}
pub fn embed_windows(&self, text: &str) -> Result<Vec<WindowEmbedding>> {
self.embed_windows_with(text, &LongTextOptions::new())
}
pub fn embed_windows_with(
&self,
text: &str,
opts: &LongTextOptions,
) -> Result<Vec<WindowEmbedding>> {
validate_long_input(text, opts)?;
let wopts = opts.window_options();
let mut chunks = chunk_long(
self.measuring_tokenizer()?,
LazyTable::new(&self.merge_table, &self.tokenizer),
text,
&wopts,
)?;
if chunks.is_empty() {
chunks.push(windit::split::Chunk::new(0, text.len()));
}
let mut windows = Vec::with_capacity(chunks.len());
let mut offset = 0usize;
for (ordinal, chunk) in chunks.iter().enumerate() {
let s = chunk.as_str(text).expect(
"chunk_long yields char-aligned boundaries (windit cuts, or 0/len from gap repair / the whole-input fallback)",
);
let ids = self.token_ids(s)?;
let embedding = self.embed_tokenized(&ids)?;
let token_span = windit::plan::Span::new(offset, ids.len(), MAX_TOKENS);
offset += ids.len();
windows.push(WindowEmbedding {
ordinal,
byte_range: chunk.start()..chunk.end(),
token_span,
embedding,
});
}
Ok(windows)
}
}
fn configure_tokenizer(tokenizer: &mut Tokenizer) -> Result<()> {
let added = tokenizer
.get_post_processor()
.map_or(0, |post| post.added_tokens(false));
if added >= MAX_TOKENS {
return Err(Error::SpecialTokenOverhead(SpecialTokenOverhead::new(
added, MAX_TOKENS,
)));
}
tokenizer
.with_truncation(Some(TruncationParams {
max_length: MAX_TOKENS,
strategy: TruncationStrategy::LongestFirst,
stride: 0,
direction: TruncationDirection::Right,
}))
.map_err(Error::TokenizerConfig)?;
tokenizer.with_padding(None);
Ok(())
}
fn validate_tokenizer_contract(tokenizer: &Tokenizer) -> Result<()> {
for (check, token, expected_id) in [
(
"special token <|startoftext|>",
contract::CLS_TOKEN,
contract::CLS_ID,
),
(
"special token <|endoftext|>",
contract::PAD_TOKEN,
contract::PAD_ID,
),
(
"special token <|return|>",
contract::EOS_TOKEN,
contract::EOS_ID,
),
] {
let actual = tokenizer.token_to_id(token);
if actual != Some(expected_id) {
return Err(Error::TokenizerContractMismatch(
TokenizerContractMismatch::new(
check,
expected_id.to_string(),
actual.map_or_else(|| "missing".to_string(), |id| id.to_string()),
),
));
}
}
let vocab_size = tokenizer.get_vocab_size(true);
if vocab_size != contract::VOCAB_SIZE {
return Err(Error::TokenizerContractMismatch(
TokenizerContractMismatch::new(
"vocab size",
contract::VOCAB_SIZE.to_string(),
vocab_size.to_string(),
),
));
}
let max_id = tokenizer.get_vocab(true).values().copied().max();
if !matches!(max_id, Some(id) if id <= contract::MAX_TOKEN_ID) {
return Err(Error::TokenizerContractMismatch(
TokenizerContractMismatch::new(
"max token id",
format!("<= {}", contract::MAX_TOKEN_ID),
max_id.map_or_else(|| "empty vocab".to_string(), |id| id.to_string()),
),
));
}
let sentinel = tokenizer
.encode(contract::SENTINEL_TEXT, true)
.map_err(Error::Tokenize)?;
if sentinel.get_ids() != contract::SENTINEL_IDS.as_slice() {
return Err(Error::TokenizerContractMismatch(
TokenizerContractMismatch::new(
"sentinel encoding",
format!("{:?}", contract::SENTINEL_IDS),
format!("{:?}", sentinel.get_ids()),
),
));
}
Ok(())
}
fn sha256_hex(bytes: &[u8]) -> String {
use sha2::{Digest, Sha256};
Sha256::digest(bytes)
.iter()
.map(|b| format!("{b:02x}"))
.collect()
}
struct Artifact {
path: std::path::PathBuf,
sha256_hex: String,
}
impl Artifact {
#[inline(always)]
const fn new(path: std::path::PathBuf, sha256_hex: String) -> Self {
Self { path, sha256_hex }
}
#[inline(always)]
fn path(&self) -> &Path {
&self.path
}
#[inline(always)]
fn sha256_hex(&self) -> &str {
&self.sha256_hex
}
}
enum TokenizerProvenance {
Artifact(Artifact),
Supplied(String),
}
fn artifact_tokenizer_path(model_path: &Path) -> std::path::PathBuf {
model_path
.parent()
.unwrap_or_else(|| Path::new(""))
.join(TOKENIZER_FILE_NAME)
}
fn validate_tokenizer_identity(provenance: &TokenizerProvenance) -> Result<()> {
let (check, sha256_hex, from) = match provenance {
TokenizerProvenance::Artifact(artifact) => (
"artifact tokenizer identity (sha-256)",
artifact.sha256_hex(),
Some(artifact.path()),
),
TokenizerProvenance::Supplied(sha256_hex) => {
("tokenizer identity (sha-256)", sha256_hex.as_str(), None)
}
};
if sha256_hex == contract::TOKENIZER_SHA256_HEX {
return Ok(());
}
Err(Error::TokenizerContractMismatch(
TokenizerContractMismatch::new(
check,
contract::TOKENIZER_SHA256_HEX.to_string(),
from.map_or_else(
|| sha256_hex.to_string(),
|path| format!("{sha256_hex} (read from {})", path.display()),
),
),
))
}
fn build_window(ids: &[u32], pad_id: i32) -> Result<([i32; MAX_TOKENS], [i32; MAX_TOKENS])> {
if ids.len() > MAX_TOKENS {
return Err(Error::TokenCount(TokenCount::new(ids.len(), MAX_TOKENS)));
}
let mut input_ids = [pad_id; MAX_TOKENS];
let mut attention_mask = [0i32; MAX_TOKENS];
for (i, &id) in ids.iter().enumerate() {
input_ids[i] = i32::try_from(id).map_err(|_| Error::TokenIdRange(id))?;
attention_mask[i] = 1;
}
Ok((input_ids, attention_mask))
}
fn validate_long_input(text: &str, opts: &LongTextOptions) -> Result<()> {
long_text::validate_long_input(text, opts, MAX_TOKENS).map_err(Error::from)
}
fn chunk_long(
measure_tok: &Tokenizer,
table: LazyTable<'_>,
text: &str,
opts: &WindowOptions,
) -> Result<Vec<windit::split::Chunk>> {
let index = TokenIndex::build(measure_tok, text)?;
let measurer = IndexMeasure::new(text, &index, measure_tok, table, opts.window());
long_text::chunk_long(&measurer, text, opts, MAX_TOKENS)
}
#[cfg(test)]
fn attach_gaps(
text: &str,
chunks: Vec<windit::split::Chunk>,
measure: &dyn Fn(usize, usize) -> Result<usize>,
window: usize,
) -> Result<Vec<windit::split::Chunk>> {
long_text::attach_gaps(text, chunks, measure, window, MAX_TOKENS)
}
#[cfg(test)]
pub(crate) fn configured_tokenizer_from_bytes(bytes: &[u8]) -> Result<Tokenizer> {
let mut tokenizer = Tokenizer::from_bytes(bytes).map_err(Error::TokenizerLoad)?;
configure_tokenizer(&mut tokenizer)?;
Ok(tokenizer)
}
#[cfg(test)]
pub(crate) fn measuring_tokenizer_from_bytes(bytes: &[u8]) -> Result<Tokenizer> {
let mut tokenizer = configured_tokenizer_from_bytes(bytes)?;
tokenizer
.with_truncation(None)
.map_err(Error::TokenizerConfig)?;
Ok(tokenizer)
}
fn describe(shape: &[usize], dtype: Option<DataType>) -> String {
let dtype = dtype.map_or("none", |d| d.as_str());
format!("{shape:?} {dtype}")
}
#[cfg(test)]
mod test_artifact;
#[cfg(test)]
mod tests;