logicaffeine_compile/model/
triage.rs

1//! What a checkpoint of an architecture no arm reads carries.
2//!
3//! A checkpoint names its architecture (`general.architecture`) and holds its tensors. When no arm of
4//! [`ModelSpec::from_source`] reads the name, the tensors still say what the model is: [`triage`] reads the checkpoint
5//! under each architecture an arm reads ([`ARCHITECTURES`]) and, for every reading that succeeds, sets the tensors the
6//! author's decode step reads for it beside the checkpoint's own. A reading whose step reads exactly the checkpoint's
7//! tensors is a family the checkpoint carries under another name; the rest are ranked by the tensors they would read
8//! that the checkpoint lacks and the checkpoint's tensors they would leave unread.
9//!
10//! Two families can read the same tensors and differ in what only the name selects — Llama's and Qwen3's rotations
11//! pair different dimensions — so a refusal names every exact reading, never a guess among them.
12
13use std::collections::BTreeSet;
14
15use super::bind::{bind, Source};
16use super::gguf::MetaValue;
17use super::program::{written, Cut};
18use super::spec::{ModelSpec, SpecError, SpecSource};
19
20/// Every architecture an arm of [`ModelSpec::from_source`] reads, in the order a tie between two readings is broken.
21pub const ARCHITECTURES: &[&str] =
22    &["llama", "qwen2", "qwen3", "qwen3moe", "qwen2moe", "qwen35", "qwen35moe", "qwen3next", "gemma3", "gpt-oss", "deepseek2", "llama4"];
23
24/// A checkpoint read under another architecture's name: its keys, still under its own, and its tensors.
25struct Named<'a, S: SpecSource + ?Sized> {
26    inner: &'a S,
27    arch: &'static str,
28}
29
30impl<S: SpecSource + ?Sized> SpecSource for Named<'_, S> {
31    fn arch(&self) -> &str {
32        self.arch
33    }
34    fn arch_meta(&self, suffix: &str) -> Option<&MetaValue> {
35        self.inner.arch_meta(suffix)
36    }
37    fn dims(&self, name: &str) -> Option<&[u64]> {
38        self.inner.dims(name)
39    }
40    fn tensor_names(&self) -> Vec<&str> {
41        self.inner.tensor_names()
42    }
43    fn has(&self, name: &str) -> bool {
44        self.inner.has(name)
45    }
46}
47
48/// One architecture's reading of a checkpoint.
49#[derive(Clone, Debug, PartialEq, Eq)]
50pub struct Reading {
51    pub arch: &'static str,
52    /// Whether the checkpoint reads as `arch` and the author writes its step; when it does not, `lacks` says why.
53    pub reads: bool,
54    /// What reading it as `arch` needs and the checkpoint lacks: a key or a tensor ([`SpecError`]'s words) or what the
55    /// author does not write yet, or the tensors the step would read that it does not hold.
56    pub lacks: Vec<String>,
57    /// The checkpoint's tensors the step would not read.
58    pub unread: Vec<String>,
59}
60
61impl Reading {
62    /// Whether the checkpoint is this architecture's under another name: every tensor read, none lacking.
63    pub fn exact(&self) -> bool {
64        self.reads && self.lacks.is_empty() && self.unread.is_empty()
65    }
66}
67
68/// The checkpoint tensors a decode step of `spec` reads: every weight its program declares, through its binding, and
69/// the rotation's frequency factors when the spec rotates by them.
70pub fn tensors_read(spec: &ModelSpec) -> BTreeSet<String> {
71    let ((_, f32s, q8s), _) = written(spec, Cut::Decode { ctx: 1 });
72    let mut read: BTreeSet<String> = f32s
73        .keys()
74        .chain(q8s.keys())
75        .filter_map(|var| match bind(spec, var) {
76            Ok(Source::Weight(b)) => Some(b.tensor),
77            _ => None,
78        })
79        .collect();
80    if super::bind::rotations(spec).iter().any(|r| r.scaling == super::spec::RopeScaling::Factors) {
81        read.insert("rope_freqs.weight".to_string());
82    }
83    read
84}
85
86/// Every architecture's reading of `g`: the exact ones first, then those that read it by how much each lacks and leaves
87/// unread, then those that do not read it.
88pub fn triage<S: SpecSource + ?Sized>(g: &S) -> Vec<Reading> {
89    let held: BTreeSet<String> = g.tensor_names().into_iter().map(str::to_string).collect();
90    let mut readings: Vec<Reading> = ARCHITECTURES
91        .iter()
92        .map(|&arch| {
93            let refused = |why: String| Reading { arch, reads: false, lacks: vec![why], unread: Vec::new() };
94            match ModelSpec::from_source(&Named { inner: g, arch }) {
95                Ok(spec) => match super::author::check(&spec) {
96                    Ok(()) => {
97                        let read = tensors_read(&spec);
98                        Reading {
99                            arch,
100                            reads: true,
101                            lacks: read.difference(&held).cloned().collect(),
102                            unread: held.difference(&read).cloned().collect(),
103                        }
104                    }
105                    Err(e) => refused(format!("the author's {}", e.0)),
106                },
107                Err(e) => refused(e.to_string()),
108            }
109        })
110        .collect();
111    readings.sort_by_key(|r| (!r.exact(), !r.reads, r.lacks.len() + r.unread.len()));
112    readings
113}
114
115/// Why `g`, of an architecture no arm reads, is refused, and what it carries: the families whose tensors it holds
116/// exactly, or the nearest reading and what it lacks and adds.
117pub fn refusal<S: SpecSource + ?Sized>(g: &S) -> String {
118    let arch = g.arch();
119    let readings = triage(g);
120    let exact: Vec<&str> = readings.iter().filter(|r| r.exact()).map(|r| r.arch).collect();
121    let head = SpecError::UnknownArchitecture(arch.to_string()).to_string();
122    if !exact.is_empty() {
123        return format!(
124            "{head}; its tensors are exactly what {} read{} — the architecture's name selects what the tensors do not \
125             (the rotation's pairing, the norms' placement)",
126            exact.iter().map(|a| format!("`{a}`")).collect::<Vec<_>>().join(" and "),
127            if exact.len() == 1 { "s" } else { "" }
128        );
129    }
130    let list = |v: &[String]| if v.is_empty() { "nothing".to_string() } else { v.join(", ") };
131    match readings.first() {
132        Some(r) if r.reads => format!(
133            "{head}; the nearest architecture an arm reads is `{}`, which lacks {} and leaves unread {}",
134            r.arch,
135            list(&r.lacks),
136            list(&r.unread)
137        ),
138        _ => format!(
139            "{head}; no architecture an arm reads takes it: {}",
140            readings.iter().map(|r| format!("`{}` needs {}", r.arch, list(&r.lacks))).collect::<Vec<_>>().join("; ")
141        ),
142    }
143}
logicaffeine_compile/model/triage.rs

logicaffeine_compile/model/
triage.rs
