Skip to main content

rigspark_core/
bootstrap.rs

1use crate::{
2    catalog::{Catalog, GgufSource, MlxSource, parse_document},
3    enrich::{Mode, RawModel, enrich},
4    sizing::{ValidationError, parse_param_count},
5};
6use serde::Deserialize;
7use std::collections::BTreeMap;
8
9pub const BOOTSTRAP_CLOCK: &str = "2026-08-04T00:00:00.000Z";
10
11#[derive(Deserialize)]
12#[serde(rename_all = "camelCase", deny_unknown_fields)]
13struct Metadata {
14    kv_bytes_per_token: Option<f64>,
15    gguf: Option<GgufSource>,
16    mlx: Option<MlxSource>,
17}
18
19pub fn build_catalog(candidates: &[RawModel], now: &str) -> Result<Catalog, ValidationError> {
20    let seed = Catalog {
21        schema_version: crate::catalog::SCHEMA_VERSION,
22        generated_at: BOOTSTRAP_CLOCK.into(),
23        models: Vec::new(),
24    };
25    let mut catalog = enrich(&seed, candidates, Mode::Backfill, now, None)?.catalog;
26    let metadata: BTreeMap<String, Metadata> =
27        parse_document(include_str!("../fixtures/bootstrap-metadata.json"))?;
28    for model in &mut catalog.models {
29        if let Some(entry) = metadata.get(&model.id) {
30            model.kv_bytes_per_token = entry.kv_bytes_per_token;
31            model.source.gguf = entry.gguf.clone();
32            model.source.mlx = entry.mlx.clone();
33        }
34        model.benchmark_proxy = family_quality_offset(&model.family)
35            .map(|_| derive_benchmark_proxy(&model.family, &model.params))
36            .transpose()?;
37    }
38    let encoded =
39        serde_json::to_string(&catalog).map_err(|error| ValidationError(error.to_string()))?;
40    Catalog::parse(&encoded)
41}
42
43pub fn family_quality_offset(family: &str) -> Option<f64> {
44    let offset = match family {
45        "kimi-k2" | "deepseek-r1" | "phi4" => 0.12,
46        "kimi-k2-thinking" => 0.14,
47        "kimi-vl" | "qwen2.5" | "gemma3n" | "gemma3" | "phi3" => 0.06,
48        "kimi-dev" | "kimi-linear" | "llama3.3" | "qwen2.5-coder" | "gemma4" | "phi3.5" => 0.08,
49        "llama3.1" | "gemma2" => 0.04,
50        "qwen3" | "deepseek-v3" | "phi4-mini" => 0.1,
51        "mixtral" | "mistral-small" | "glm4" => 0.05,
52        "mistral-nemo" | "yi" | "granite3.1" => 0.03,
53        "mistral" | "olmo2" | "granite3-moe" => 0.02,
54        "llama3.2" | "smollm2" => 0.0,
55        _ => return None,
56    };
57    Some(offset)
58}
59
60pub fn derive_benchmark_proxy(family: &str, params: &str) -> Result<f64, ValidationError> {
61    let offset = family_quality_offset(family).unwrap_or(0.0);
62    let billions = parse_param_count(params)? / 1e9;
63    let base = [
64        (0.5, 0.05),
65        (1.0, 0.12),
66        (2.0, 0.2),
67        (4.0, 0.28),
68        (7.0, 0.36),
69        (10.0, 0.44),
70        (15.0, 0.52),
71        (25.0, 0.58),
72        (40.0, 0.64),
73        (80.0, 0.7),
74        (200.0, 0.78),
75    ]
76    .into_iter()
77    .find_map(|(ceiling, score)| (billions < ceiling).then_some(score))
78    .unwrap_or(0.85_f64);
79    Ok(((base + offset).clamp(0.05, 0.99) * 100.0).round() / 100.0)
80}