candding 0.1.0

//! Construction of a `TextEmbedding` from a model id or directory.

use std::fs;

use candle_core::DType;
use candle_core::Device;
pub(crate) use dtype::reject_unsupported_dtype;
pub(crate) use dtype::resolve_dtype;

use super::fit;
use super::text::TextEmbedding;
use crate::device;
use crate::error::CanddingError;
use crate::error::Result;
use crate::hub;
use crate::hub::ModelSource;
use crate::layers::dense::Dense;
use crate::loader;
use crate::loader::ModelFiles;
use crate::models;
use crate::models::Encoder;
use crate::pooling::Pooling;
use crate::registry;
use crate::registry::Family;
use crate::registry::ModelDescriptor;
use crate::tokenizer::TextTokenizer;
use crate::tokenizer::read_pad_token;
use crate::weights;

mod dtype;

pub struct TextEmbeddingBuilder {
  spec: String,
  revision: String,
  device: Option<Device>,
  dtype: Option<DType>,
  max_length: Option<usize>,
  batch_size: usize,
  descriptor: Option<ModelDescriptor>,
  normalize: Option<bool>,
}

impl TextEmbeddingBuilder {
  pub(crate) fn new(spec: impl Into<String>) -> Self {
    Self {
      spec: spec.into(),
      revision: "main".to_string(),
      device: None,
      dtype: None,
      max_length: None,
      batch_size: 32,
      descriptor: None,
      normalize: None,
    }
  }

  pub fn device(mut self, device: Device) -> Self {
    self.device = Some(device);
    self
  }

  pub fn dtype(mut self, dtype: DType) -> Self {
    self.dtype = Some(dtype);
    self
  }

  /// Capped at the model's `max_position_embeddings` when building.
  pub fn max_length(mut self, max_length: usize) -> Self {
    self.max_length = Some(max_length);
    self
  }

  pub fn batch_size(mut self, batch_size: usize) -> Self {
    self.batch_size = batch_size.max(1);
    self
  }

  pub fn revision(mut self, revision: impl Into<String>) -> Self {
    self.revision = revision.into();
    self
  }

  /// Supplies the descriptor for a repository or directory outside the registry, or replaces a dense entry's own;
  /// an id only another catalog carries is refused with or without one.
  pub fn descriptor(mut self, descriptor: ModelDescriptor) -> Self {
    self.descriptor = Some(descriptor);
    self
  }

  /// Overrides the descriptor's `normalize` flag, which is `true` for every model, registered or not, whatever its
  /// `modules.json` lists. The model-independent way to get the raw pooled vector back.
  pub fn normalize(mut self, normalize: bool) -> Self {
    self.normalize = Some(normalize);
    self
  }

  /// Loads the model. An id only another catalog carries, and a dtype the descriptor refuses on the device, fail
  /// before anything is fetched; a local directory or an id no catalog carries loads as an unregistered snapshot.
  pub fn build(self) -> Result<TextEmbedding> {
    let source = canonical_source(&self.spec, &self.revision);
    refuse_another_catalogs_id(&self.spec, &source)?;
    let registered = self
      .descriptor
      .or_else(|| registry::dense::lookup(&self.spec));
    let device = match self.device {
      Some(device) => device,
      None => device::auto()?,
    };
    let default_dtype = registered
      .as_ref()
      .map(|d| d.default_dtype)
      .unwrap_or(DType::F32);
    let dtype = resolve_dtype(&self.spec, self.dtype, default_dtype, &device)?;
    if let Some(descriptor) = registered.as_ref() {
      reject_unsupported_dtype(&self.spec, dtype, descriptor.dtype_backends, &device)?;
      let (size_gb, stored) = (descriptor.size_gb, descriptor.default_dtype);
      fit::check_download(&self.spec, &device, &source, size_gb, stored, dtype)?;
    }
    let dir = hub::resolve(&source, registry::weights_format(&self.spec))?;
    let files = loader::discover(&dir, loader::Pipeline::Dense)?;
    let mut weights = fit::at(&files.weights, dtype);
    for dense_dir in &files.dense_dirs {
      weights.push((dense_dir.join("model.safetensors"), DType::F32));
    }
    fit::check_files(&self.spec, &device, &weights)?;
    let config_json = fs::read_to_string(&files.config)?;
    let architecture = loader::read_architecture(&files.config)?;
    let family = Family::from_model_type(&architecture)
      .ok_or_else(|| CanddingError::UnsupportedArchitecture(architecture.clone()))?;
    tracing::info!(model = %self.spec, device = device::describe(&device), dtype = dtype.as_str(), "loading");
    let vb = unsafe { weights::var_builder(&files.weights, dtype, &device)? };
    let encoder = models::load_encoder(&architecture, &config_json, vb)?;
    let dense = files
      .dense_dirs
      .iter()
      .map(|dir| Dense::load(dir, &device))
      .collect::<Result<Vec<_>>>()?;
    let mut descriptor = match registered {
      Some(descriptor) => descriptor,
      None => descriptor_from_snapshot(&self.spec, family, &files, encoder.as_ref(), &dense)?,
    };
    let found = dense_names(&files);
    if descriptor.dense_dirs != found {
      return Err(CanddingError::DenseMismatch {
        expected: descriptor.dense_dirs.clone(),
        found,
      });
    }
    if let Some(normalize) = self.normalize {
      descriptor.normalize = normalize;
    }
    let max_length = self
      .max_length
      .unwrap_or(descriptor.default_max_length)
      .min(encoder.max_position_embeddings());
    let pad = read_pad_token(&config_json, files.tokenizer_config.as_deref())?;
    let tokenizer = TextTokenizer::from_file(&files.tokenizer, max_length, pad)?;
    Ok(TextEmbedding::new(
      encoder,
      tokenizer,
      dense,
      descriptor,
      device,
      dtype,
      self.batch_size,
    ))
  }
}

/// `registry::other_catalog`'s answer for a hub id as the builder's refusal, naming the builder and command that
/// load the id. A local directory is a path rather than a registry id, whatever it is named, so it is not refused.
fn refuse_another_catalogs_id(spec: &str, source: &ModelSource) -> Result<()> {
  if let ModelSource::Local(_) = source {
    return Ok(());
  }
  match registry::other_catalog(spec) {
    Some(catalog) => Err(CanddingError::NonDenseModel {
      model: spec.to_string(),
      catalog,
    }),
    None => Ok(()),
  }
}

/// Where a spec is resolved from, under the registry's spelling of any id it carries: the cache is case-sensitive,
/// so a differently cased id would otherwise fetch a second copy. A local directory keeps the caller's spelling.
fn canonical_source(spec: &str, revision: &str) -> ModelSource {
  match ModelSource::parse(spec, revision) {
    ModelSource::Hub { repo_id, revision } => {
      let canonical = registry::canonical_id(&repo_id);
      ModelSource::Hub {
        repo_id: canonical.map_or(repo_id, str::to_string),
        revision,
      }
    }
    local => local,
  }
}

/// The directory names (`2_Dense`, `3_Dense`) the loader found, for the descriptor check and snapshot descriptors.
fn dense_names(files: &ModelFiles) -> Vec<String> {
  files
    .dense_dirs
    .iter()
    .filter_map(|p| p.file_name())
    .map(|n| n.to_string_lossy().into_owned())
    .collect()
}

/// Unregistered repositories follow their sentence-transformers files for pooling, Dense modules and the default
/// length; their rows are unit length whatever `modules.json` lists, as every catalog model's are.
fn descriptor_from_snapshot(
  id: &str,
  family: Family,
  files: &ModelFiles,
  encoder: &dyn Encoder,
  dense: &[Dense],
) -> Result<ModelDescriptor> {
  let pooling = match &files.pooling_config {
    Some(path) => loader::read_pooling(path)?,
    None => Pooling::Mean,
  };
  let max_seq_length = match &files.sentence_bert_config {
    Some(path) => loader::read_max_seq_length(path)?,
    None => None,
  };
  let dim = dense
    .last()
    .map(Dense::out_features)
    .unwrap_or(encoder.hidden_size());
  let mut descriptor = ModelDescriptor::new(id, family, dim, pooling);
  descriptor.default_max_length = max_seq_length
    .unwrap_or(descriptor.default_max_length)
    .min(encoder.max_position_embeddings());
  descriptor.dense_dirs = dense_names(files);
  Ok(descriptor)
}

#[cfg(test)]
mod tests;