candding 0.1.0

//! Discovery of the files inside a model snapshot.

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::path::PathBuf;

use serde::Deserialize;

use crate::error::CanddingError;
use crate::error::Result;
use crate::pooling::Pooling;

#[derive(Clone, Debug)]
pub struct ModelFiles {
  pub dir: PathBuf,
  pub config: PathBuf,
  pub tokenizer: PathBuf,
  pub tokenizer_config: Option<PathBuf>,
  pub weights: Vec<PathBuf>,
  pub pooling_config: Option<PathBuf>,
  pub sentence_bert_config: Option<PathBuf>,
  pub modules: Option<PathBuf>,
  pub dense_dirs: Vec<PathBuf>,
  /// `sparse_linear.pt`, bge-m3's sparse head, present only when the snapshot was fetched with it.
  pub sparse_linear: Option<PathBuf>,
  /// `colbert_linear.pt`, bge-m3's multi-vector head, resolved the same optional way.
  pub colbert_linear: Option<PathBuf>,
}

/// bge-m3's sparse head file, shared by `discover` and `models::load_sparse`'s error so the two cannot drift.
pub(crate) const SPARSE_LINEAR_FILENAME: &str = "sparse_linear.pt";

/// bge-m3's multi-vector head file, shared by `discover` and `models::load_multi_vector`'s error.
pub(crate) const COLBERT_LINEAR_FILENAME: &str = "colbert_linear.pt";

/// `Qdrant/minicoil-v1`'s three files, resolved in that repository's own snapshot, not beside the encoder.
pub(crate) const MINICOIL_MATRIX_FILENAME: &str = "minicoil.triplet.model.npy";
pub(crate) const MINICOIL_VOCAB_FILENAME: &str = "minicoil.triplet.model.vocab";
pub(crate) const MINICOIL_STOPWORDS_FILENAME: &str = "stopwords.txt";

/// `Qdrant/minicoil-v1`'s own three files, already resolved to real paths on disk.
#[derive(Clone, Debug)]
pub(crate) struct MiniCoilSideFiles {
  pub matrix: PathBuf,
  pub vocab: PathBuf,
  pub stopwords: PathBuf,
}

/// Resolves MiniCOIL's three files in `head_source`'s snapshot; all are required, unlike the optional `sparse_linear`.
pub(crate) fn discover_minicoil_files(head_source_dir: &Path) -> Result<MiniCoilSideFiles> {
  Ok(MiniCoilSideFiles {
    matrix: required(head_source_dir, MINICOIL_MATRIX_FILENAME)?,
    vocab: required(head_source_dir, MINICOIL_VOCAB_FILENAME)?,
    stopwords: required(head_source_dir, MINICOIL_STOPWORDS_FILENAME)?,
  })
}

/// Which pipeline asks `discover` to validate `modules.json`: Splade_PP publishes two module types only the
/// sparse pipeline runs, which the dense pipeline must still refuse.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Pipeline {
  Dense,
  Sparse,
  /// `embed::rerank::TextCrossEncoderBuilder::build`: runs `RERANK_MODULE_TYPES` and refuses every other type,
  /// since no rerank head runs a `Pooling`, `Dense` or `Normalize` module.
  Rerank,
}

pub fn discover(dir: &Path, pipeline: Pipeline) -> Result<ModelFiles> {
  if !dir.is_dir() {
    return Err(CanddingError::MissingDir(dir.to_path_buf()));
  }
  let modules = optional(dir, "modules.json");
  let dense_dirs = match &modules {
    Some(path) => dense_dirs(dir, &read_modules(path)?, pipeline)?,
    None => Vec::new(),
  };
  Ok(ModelFiles {
    dir: dir.to_path_buf(),
    config: required(dir, "config.json")?,
    tokenizer: required(dir, "tokenizer.json")?,
    tokenizer_config: optional(dir, "tokenizer_config.json"),
    weights: weight_files(dir)?,
    pooling_config: optional(dir, "1_Pooling/config.json"),
    sentence_bert_config: optional(dir, "sentence_bert_config.json"),
    modules,
    dense_dirs,
    sparse_linear: optional(dir, SPARSE_LINEAR_FILENAME),
    colbert_linear: optional(dir, COLBERT_LINEAR_FILENAME),
  })
}

/// The `modules.json` types the dense pipeline runs: three need no files, `Dense` loads here. The sparse
/// pipeline accepts these too (bge-m3 reaches it with this list), plus `SPARSE_MODULE_TYPES`.
const DENSE_MODULE_TYPES: [&str; 4] = ["Transformer", "Pooling", "Normalize", "Dense"];

/// The `modules.json` types a reranker's pipeline is made of: the transformer, and the `LogitScore` head the Qwen3
/// rerankers publish as `sentence_transformers.cross_encoder.modules.logit_score.LogitScore`.
const RERANK_MODULE_TYPES: [&str; 2] = ["Transformer", "LogitScore"];

/// The two types only the sparse pipeline runs, as Splade_PP publishes them, matched by the name after the last `.`.
const SPARSE_MODULE_TYPES: [&str; 2] = ["MLMTransformer", "SpladePooling"];

/// `{n}_Dense` directories named by `modules.json`, in pipeline order. A type the asking pipeline does not run is
/// an error, never skipped, as it changes the vectors; the dense pipeline names a sparse-only one `SparseOnlyModule`.
fn dense_dirs(dir: &Path, modules: &[Module], pipeline: Pipeline) -> Result<Vec<PathBuf>> {
  for module in modules {
    let kind = module.kind.rsplit('.').next().unwrap_or_default();
    if pipeline == Pipeline::Rerank {
      if RERANK_MODULE_TYPES.contains(&kind) {
        continue;
      }
      return Err(CanddingError::UnsupportedModule(module.kind.clone()));
    }
    if DENSE_MODULE_TYPES.contains(&kind) {
      continue;
    }
    if SPARSE_MODULE_TYPES.contains(&kind) {
      if pipeline == Pipeline::Sparse {
        continue;
      }
      return Err(CanddingError::SparseOnlyModule(module.kind.clone()));
    }
    return Err(CanddingError::UnsupportedModule(module.kind.clone()));
  }
  modules
    .iter()
    .filter(|m| m.kind.ends_with(".Dense"))
    .map(|m| {
      let path = dir.join(&m.path);
      required(&path, "config.json")?;
      required(&path, "model.safetensors")?;
      Ok(path)
    })
    .collect()
}

fn required(dir: &Path, name: &str) -> Result<PathBuf> {
  let path = dir.join(name);
  path
    .is_file()
    .then_some(path)
    .ok_or_else(|| CanddingError::MissingFile(dir.join(name)))
}

fn optional(dir: &Path, name: &str) -> Option<PathBuf> {
  let path = dir.join(name);
  path.is_file().then_some(path)
}

#[derive(Deserialize)]
struct WeightIndex {
  weight_map: HashMap<String, String>,
}

/// `model.safetensors`, else every shard named in `model.safetensors.index.json`, else `pytorch_model.bin` for a
/// repository that publishes only a torch pickle, as `BAAI/bge-m3` does.
fn weight_files(dir: &Path) -> Result<Vec<PathBuf>> {
  let single = dir.join("model.safetensors");
  if single.is_file() {
    return Ok(vec![single]);
  }
  let index = dir.join("model.safetensors.index.json");
  if index.is_file() {
    let parsed: WeightIndex = serde_json::from_str(&fs::read_to_string(&index)?)?;
    let shards: BTreeSet<String> = parsed.weight_map.into_values().collect();
    return Ok(shards.into_iter().map(|shard| dir.join(shard)).collect());
  }
  let pickle = dir.join("pytorch_model.bin");
  if pickle.is_file() {
    return Ok(vec![pickle]);
  }
  Err(CanddingError::MissingWeights(dir.to_path_buf()))
}

#[derive(Deserialize)]
struct RawConfig {
  model_type: String,
  #[serde(default)]
  position_embedding_type: Option<String>,
}

/// The architecture key a checkpoint dispatches on: `config.json`'s `model_type`, except `jina_bert` for a `bert`
/// with `"position_embedding_type": "alibi"`. The single resolution point for every dispatch.
pub fn read_architecture(config: &Path) -> Result<String> {
  let raw: RawConfig = serde_json::from_str(&fs::read_to_string(config)?)?;
  if raw.model_type == "bert" && raw.position_embedding_type.as_deref() == Some("alibi") {
    return Ok("jina_bert".to_string());
  }
  Ok(raw.model_type)
}

/// The `1_Pooling/config.json` modes candding implements; sentence-transformers also has `max_tokens`,
/// `weightedmean_tokens` and `mean_sqrt_len_tokens`.
const POOLING_MODES: [&str; 3] = [
  "pooling_mode_cls_token",
  "pooling_mode_mean_tokens",
  "pooling_mode_lasttoken",
];

/// sentence-transformers `1_Pooling/config.json`; no mode at all means mean, the library's default.
/// An unimplemented mode is an error rather than a mean, since each returns a different vector.
pub fn read_pooling(path: &Path) -> Result<Pooling> {
  let raw: BTreeMap<String, serde_json::Value> = serde_json::from_str(&fs::read_to_string(path)?)?;
  let requested = |key: &str| raw.get(key).and_then(serde_json::Value::as_bool) == Some(true);
  let unknown = raw.iter().find(|(key, value)| {
    key.starts_with("pooling_mode_")
      && value.as_bool() == Some(true)
      && !POOLING_MODES.contains(&key.as_str())
  });
  if let Some((mode, _)) = unknown {
    return Err(CanddingError::UnsupportedPooling(mode.clone()));
  }
  Ok(if requested("pooling_mode_cls_token") {
    Pooling::Cls
  } else if requested("pooling_mode_lasttoken") {
    Pooling::LastToken
  } else {
    Pooling::Mean
  })
}

#[derive(Default, Deserialize)]
#[serde(default)]
struct SentenceBertConfig {
  max_seq_length: Option<usize>,
}

/// The trained sequence length from `sentence_bert_config.json`, often below the position table (256 against 512
/// for all-MiniLM-L6-v2); `None` when the key is missing or zero.
pub fn read_max_seq_length(path: &Path) -> Result<Option<usize>> {
  let cfg: SentenceBertConfig = serde_json::from_str(&fs::read_to_string(path)?)?;
  Ok(cfg.max_seq_length.filter(|&len| len > 0))
}

#[derive(Clone, Debug, Deserialize, PartialEq, Eq)]
pub struct Module {
  #[serde(rename = "type")]
  pub kind: String,
  #[serde(default)]
  pub path: String,
}

/// Pipeline entries from `modules.json` in order.
pub fn read_modules(path: &Path) -> Result<Vec<Module>> {
  Ok(serde_json::from_str(&fs::read_to_string(path)?)?)
}

#[cfg(test)]
mod tests;