candding 0.1.0

use candle_core::DType;

use super::DtypeSupport;
use super::Family;
use super::Support;
use super::WeightsFormat;
use super::entry::Entry;
use super::t4::T4_BF16_UNTESTED;
use crate::pooling::Pooling;

const BGE_EN_QUERY: &str = "Represent this sentence for searching relevant passages: {text}";
const BGE_ZH_QUERY: &str = "为这个句子生成表示以用于检索相关文章：{text}";

/// A BERT-family dense entry with the statuses all of them share: every dtype verified on the CPU and Metal that
/// candle can run there, and F32 and F16 on the T4.
const fn bert(
  id: &'static str,
  dim: usize,
  default_max_length: usize,
  pooling: Pooling,
  query_template: Option<&'static str>,
  license: &'static str,
  size_gb: f32,
) -> Entry {
  Entry {
    id,
    family: Family::Bert,
    dim,
    default_max_length,
    pooling,
    cpu: Support::Verified,
    metal: Support::Verified,
    cuda: Support::Verified,
    query_template,
    passage_template: None,
    default_instruction: None,
    dense_dirs: &[],
    mrl: false,
    layer_norm_after_pooling: false,
    gated: false,
    default_dtype: DType::F32,
    dtypes: DtypeSupport::verified(),
    license,
    size_gb,
    phase: 0,
    weights: WeightsFormat::Safetensors,
    accel_only_dtypes: &[],
    cell_exceptions: T4_BF16_UNTESTED,
  }
}

pub(super) const ENTRIES: &[Entry] = &[
  bert(
    "BAAI/bge-small-en-v1.5",
    384,
    512,
    Pooling::Cls,
    Some(BGE_EN_QUERY),
    "MIT",
    0.13,
  ),
  bert(
    "BAAI/bge-base-en-v1.5",
    768,
    512,
    Pooling::Cls,
    Some(BGE_EN_QUERY),
    "MIT",
    0.44,
  ),
  bert(
    "BAAI/bge-large-en-v1.5",
    1024,
    512,
    Pooling::Cls,
    Some(BGE_EN_QUERY),
    "MIT",
    1.34,
  ),
  bert(
    "BAAI/bge-small-zh-v1.5",
    512,
    512,
    Pooling::Cls,
    Some(BGE_ZH_QUERY),
    "MIT",
    0.10,
  ),
  bert(
    "sentence-transformers/all-MiniLM-L6-v2",
    384,
    256,
    Pooling::Mean,
    None,
    "Apache-2.0",
    0.09,
  ),
  bert(
    "thenlper/gte-base",
    768,
    512,
    Pooling::Mean,
    None,
    "MIT",
    0.22,
  ),
  bert(
    "thenlper/gte-large",
    1024,
    512,
    Pooling::Mean,
    None,
    "MIT",
    0.67,
  ),
  bert(
    "mixedbread-ai/mxbai-embed-large-v1",
    1024,
    512,
    Pooling::Cls,
    Some(BGE_EN_QUERY),
    "Apache-2.0",
    0.67,
  ),
  bert(
    "Snowflake/snowflake-arctic-embed-xs",
    384,
    512,
    Pooling::Cls,
    Some(BGE_EN_QUERY),
    "Apache-2.0",
    0.09,
  ),
  bert(
    "Snowflake/snowflake-arctic-embed-s",
    384,
    512,
    Pooling::Cls,
    Some(BGE_EN_QUERY),
    "Apache-2.0",
    0.13,
  ),
  bert(
    "Snowflake/snowflake-arctic-embed-m",
    768,
    512,
    Pooling::Cls,
    Some(BGE_EN_QUERY),
    "Apache-2.0",
    0.44,
  ),
  bert(
    "Snowflake/snowflake-arctic-embed-l",
    1024,
    512,
    Pooling::Cls,
    Some(BGE_EN_QUERY),
    "Apache-2.0",
    1.34,
  ),
  bert(
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    384,
    128,
    Pooling::Mean,
    None,
    "Apache-2.0",
    0.47,
  ),
  E5_SMALL,
];

/// `intfloat/multilingual-e5-small` is a BERT by its `config.json`: a multilingual MiniLM distilled from XLM-R, with
/// the XLM-R tokenizer, two token types, 512 absolute positions and an epsilon of `1e-12`; it pads with id 0, `<s>`.
const E5_SMALL: Entry = Entry {
  id: "intfloat/multilingual-e5-small",
  family: Family::Bert,
  dim: 384,
  default_max_length: 512,
  pooling: Pooling::Mean,
  cpu: Support::Verified,
  metal: Support::Verified,
  cuda: Support::Verified,
  query_template: Some("query: {text}"),
  passage_template: Some("passage: {text}"),
  default_instruction: None,
  dense_dirs: &[],
  mrl: false,
  layer_norm_after_pooling: false,
  gated: false,
  default_dtype: DType::F32,
  dtypes: DtypeSupport::verified(),
  license: "MIT",
  size_gb: 0.47,
  phase: 2,
  weights: WeightsFormat::Safetensors,
  accel_only_dtypes: &[],
  cell_exceptions: T4_BF16_UNTESTED,
};