candding 0.1.0

use std::fs;

use candle_core::DType;
use candle_core::Device;

use super::MultiVectorTextEmbedding;
use super::build;
use crate::device;
use crate::embed::builder::reject_unsupported_dtype;
use crate::embed::builder::resolve_dtype;
use crate::embed::fit;
use crate::error::CanddingError;
use crate::error::Result;
use crate::hub;
use crate::hub::ModelSource;
use crate::loader;
use crate::models;
use crate::registry;
use crate::registry::Catalog;
use crate::registry::multi_vector::MultiVectorDescriptor;
use crate::tokenizer::TextTokenizer;
use crate::tokenizer::read_pad_token;
use crate::weights;

/// Eight, as in the other pipelines; a multi-vector head has no vocabulary-sized intermediate to size it against.
const DEFAULT_BATCH_SIZE: usize = 8;

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

impl MultiVectorTextEmbeddingBuilder {
  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: DEFAULT_BATCH_SIZE,
      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
  }

  /// The document-side length, capped at the model's `max_position_embeddings`; a descriptor's fixed `query_length`
  /// keeps the query width regardless.
  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 or overrides the descriptor, for a multi-vector head outside the registry.
  pub fn descriptor(mut self, descriptor: MultiVectorDescriptor) -> Self {
    self.descriptor = Some(descriptor);
    self
  }

  /// Overrides the descriptor's `normalize`, `true` for every registered model: `false` returns each row as the
  /// projection writes it, the same model-independent switch `TextEmbeddingBuilder::normalize` is.
  pub fn normalize(mut self, normalize: bool) -> Self {
    self.normalize = Some(normalize);
    self
  }

  pub fn build(self) -> Result<MultiVectorTextEmbedding> {
    let overridden = self.descriptor.is_some();
    let mut descriptor = match self.descriptor {
      Some(descriptor) => descriptor,
      None => registry::multi_vector::lookup(&self.spec).ok_or_else(|| {
        CanddingError::UnregisteredModel {
          id: self.spec.clone(),
          catalog: Catalog::MultiVector,
        }
      })?,
    };
    if let Some(normalize) = self.normalize {
      descriptor.normalize = normalize;
    }
    let spec = if overridden {
      self.spec.as_str()
    } else {
      descriptor.id.as_str()
    };
    build::require_fixed_query_length(spec, &descriptor)?;
    let device = match self.device {
      Some(device) => device,
      None => device::auto()?,
    };
    let dtype = resolve_dtype(&self.spec, self.dtype, DType::F32, &device)?;
    reject_unsupported_dtype(&self.spec, dtype, descriptor.dtype_backends, &device)?;
    let source = ModelSource::parse(spec, &self.revision);
    let size_gb = descriptor.size_gb;
    fit::check_download(&self.spec, &device, &source, size_gb, DType::F32, dtype)?;
    let dir = hub::resolve(&source, descriptor.weights)?;
    let files = loader::discover(&dir, loader::Pipeline::Dense)?;
    let weight_files = build::weight_files(&descriptor, &files, dtype);
    fit::check_files(&self.spec, &device, &weight_files)?;
    let config_json = fs::read_to_string(&files.config)?;
    tracing::info!(model = %self.spec, device = device::describe(&device), dtype = dtype.as_str(), "loading multi-vector model");
    let vb = unsafe { weights::var_builder(&files.weights, dtype, &device)? };
    let model = models::load_multi_vector(
      descriptor.head,
      &config_json,
      vb,
      files.colbert_linear.as_deref(),
      &files.tokenizer,
      build::conventions_from_descriptor(&descriptor),
    )?;
    let max_length = self
      .max_length
      .unwrap_or(descriptor.default_max_length)
      .min(model.max_position_embeddings());
    let pad = read_pad_token(&config_json, files.tokenizer_config.as_deref())?;
    let mut tokenizer = TextTokenizer::from_file(&files.tokenizer, max_length, pad.clone())?;
    build::register_marker(&mut tokenizer, descriptor.document_marker.as_deref())?;
    let query_tokenizer = build::build_query_tokenizer(
      &files.tokenizer,
      max_length,
      pad,
      &descriptor,
      model.max_position_embeddings(),
      files.tokenizer_config.as_deref(),
    )?;
    Ok(MultiVectorTextEmbedding::new(
      model,
      tokenizer,
      query_tokenizer,
      descriptor,
      device,
      dtype,
      self.batch_size,
    ))
  }
}

#[cfg(test)]
mod tests;