candding 0.1.0

//! Resolving a checkpoint's special-token names from its published config, not a shared literal.

use std::fs;
use std::path::Path;

use serde::Deserialize;

use crate::error::Result;

/// Where the padding id comes from, by priority: `config.json`, `tokenizer_config.json`'s `pad_token`, then 0.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum PadToken {
  Id(u32),
  Named(String),
  Unknown,
}

/// Read as an `i64` because `config.json` writes `-1` for a model that has no pad token, and a `u32` field would
/// fail the whole build on that rather than reading on.
#[derive(Deserialize)]
struct PadConfig {
  #[serde(default)]
  pad_token_id: Option<i64>,
}

#[derive(Deserialize)]
struct TokenizerConfig {
  #[serde(default)]
  pad_token: Option<PadTokenField>,
  #[serde(default)]
  mask_token: Option<PadTokenField>,
}

#[derive(Deserialize)]
#[serde(untagged)]
enum PadTokenField {
  Plain(String),
  Added { content: String },
}

/// `config.json`'s `pad_token_id`, then `tokenizer_config.json`'s `pad_token`, then `PadToken::Unknown`; an id
/// outside `u32` such as `-1` counts as absent.
pub fn read_pad_token(config_json: &str, tokenizer_config: Option<&Path>) -> Result<PadToken> {
  let cfg: PadConfig = serde_json::from_str(config_json)?;
  if let Some(id) = cfg.pad_token_id.and_then(|id| u32::try_from(id).ok()) {
    return Ok(PadToken::Id(id));
  }
  if let Some(path) = tokenizer_config {
    let tc: TokenizerConfig = serde_json::from_str(&fs::read_to_string(path)?)?;
    if let Some(field) = tc.pad_token {
      return Ok(PadToken::Named(match field {
        PadTokenField::Plain(name) => name,
        PadTokenField::Added { content } => content,
      }));
    }
  }
  Ok(PadToken::Unknown)
}

/// The fill token of a fixed-length query, from `tokenizer_config.json`'s `mask_token` (`<mask>` for
/// jina-colbert-v2), falling back to `[MASK]`.
pub fn read_mask_token(tokenizer_config: Option<&Path>) -> Result<String> {
  if let Some(path) = tokenizer_config {
    let tc: TokenizerConfig = serde_json::from_str(&fs::read_to_string(path)?)?;
    if let Some(field) = tc.mask_token {
      return Ok(match field {
        PadTokenField::Plain(name) => name,
        PadTokenField::Added { content } => content,
      });
    }
  }
  Ok("[MASK]".to_string())
}