topk_probe/
topk_probe.rs

1//! Next-token top-k probe over a RAW token prefix (scratch diagnostic).
2//!
3//! `explain` applies the chat template and only ever reports position 0,
4//! so it cannot answer "what does the model think comes after token N of
5//! its own answer". This feeds an arbitrary raw string (added tokens are
6//! honoured) and prints the top-k logits AND softmax mass of the next
7//! token — the margin is what separates "quantization flipped a near-tie"
8//! from "the model really believes this".
9//!
10//! Usage:
11//!   cargo run --release --example topk_probe -- <model.cmf> <raw text> [k]
12
13use cortiq_core::CmfModel;
14use cortiq_engine::{Pipeline, SamplerConfig};
15use std::sync::Arc;
16
17fn main() {
18    let args: Vec<String> = std::env::args().collect();
19    if args.len() < 3 {
20        eprintln!("usage: topk_probe <model.cmf> <raw text> [k]");
21        std::process::exit(2);
22    }
23    let k: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(10);
24
25    let model = Arc::new(CmfModel::open_sharded(&args[1]).expect("open model"));
26    let mut pipeline = Pipeline::from_model(&model, SamplerConfig::default()).expect("pipeline");
27    let ids = pipeline.tokenizer.encode(&args[2]);
28    eprintln!("prompt tokens: {}", ids.len());
29    eprintln!("last 8 ids: {:?}", &ids[ids.len().saturating_sub(8)..]);
30
31    let logits = pipeline.prefill_next_logits(&ids, None);
32    let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
33    let sum: f64 = logits.iter().map(|&l| ((l - max) as f64).exp()).sum();
34
35    let mut order: Vec<usize> = (0..logits.len()).collect();
36    order.sort_by(|&a, &b| logits[b].total_cmp(&logits[a]));
37    for &i in order.iter().take(k) {
38        let p = ((logits[i] - max) as f64).exp() / sum;
39        println!(
40            "{:>8}  logit {:>9.4}  p {:.6}  {:?}",
41            i,
42            logits[i],
43            p,
44            pipeline.tokenizer.decode_token(i as u32)
45        );
46    }
47}
topk_probe/topk_probe.rs

topk_probe/
topk_probe.rs
