lm_forward_probe/
lm_forward_probe.rs

1//! Probe: run the LM forward on the real prompt and inspect next-token
2//! logits (top-10), then compare greedy vs sampled code generation.
3//!
4//! Usage: cargo run --release --example lm_forward_probe -- <model_dir>
5
6use burn::tensor::{Int, Tensor, TensorData};
7use maolan_generate::acestep::lm::{self, AceStepLm, AudioCodeVocab, SamplingConfig};
8use maolan_generate::acestep::qwen3::{Qwen3Config, Qwen3Model};
9use std::path::Path;
10
11type B = burn::backend::NdArray<f32>;
12
13fn main() -> anyhow::Result<()> {
14    let model_dir = std::env::args()
15        .nth(1)
16        .unwrap_or_else(|| "/home/meka/repos/ace".to_string());
17    let model_dir = Path::new(&model_dir);
18    let device = Default::default();
19
20    let config = Qwen3Config::load(&model_dir.join("lm_config.json"))?;
21    let lm = AceStepLm::<B>::from_burnpack(&config, &model_dir.join("acestep-lm.bpk"), &device)?;
22    let vocab = AudioCodeVocab::from_tokenizer_json(&model_dir.join("lm_tokenizer.json"))?;
23
24    let cot = lm::build_cot_block(
25        "Metal guitar with a lot of distortion",
26        Some(120.0),
27        Some("A minor"),
28        Some("4/4"),
29        4,
30    );
31    let prompt = lm::build_codes_prompt("Metal guitar with a lot of distortion", &cot);
32    let ids = lm::tokenize_prompt(&model_dir.join("lm_tokenizer.json"), &prompt)?;
33    println!("prompt: {} tokens", ids.len());
34
35    // One full forward; inspect logits at the last position.
36    let model: &Qwen3Model<B> = &lm.model;
37    let ids_i64: Vec<i64> = ids.iter().map(|&id| i64::from(id)).collect();
38    let len = ids_i64.len();
39    let input = Tensor::<B, 2, Int>::from_data(TensorData::new(ids_i64, [1, len]), &device);
40    let hidden = model.forward(input, true);
41    let [_, _, hidden_dim] = hidden.dims();
42    let last = hidden.narrow(1, len - 1, 1).reshape([1, hidden_dim]);
43    let weight = model.embedding_weight();
44    let logits = last
45        .matmul(weight.clone().transpose())
46        .reshape([config.vocab_size as usize]);
47    let values: Vec<f32> = logits
48        .into_data()
49        .convert::<f32>()
50        .to_vec()
51        .map_err(|e| anyhow::anyhow!("{e}"))?;
52    let mut order: Vec<usize> = (0..values.len()).collect();
53    order.sort_by(|&a, &b| values[b].total_cmp(&values[a]));
54    println!("\ntop-10 next-token predictions:");
55    for &idx in order.iter().take(10) {
56        let id = idx as u32;
57        let label = vocab
58            .token_id_to_code(id)
59            .map(|c| format!("audio_code_{c}"))
60            .unwrap_or_else(|| format!("token {id}"));
61        println!("  id {id:>7} logit {:>8.3}  {label}", values[idx]);
62    }
63
64    // Greedy generation (temperature ~0) vs sampled (official defaults).
65    let greedy = SamplingConfig {
66        temperature: 1e-5,
67        top_k: 1,
68        cfg_scale: 1.0,
69        ..SamplingConfig::new(28, 0)
70    };
71    for (name, sampling) in [
72        ("greedy", greedy),
73        ("sampled official", SamplingConfig::new(28, 0)),
74    ] {
75        let codes = lm.generate_codes(&ids, None, &vocab, &sampling, None);
76        println!("\n{name}: {} codes: {codes:?}", codes.len());
77    }
78
79    // Greedy generation via FULL forward each step (no KV cache) to isolate
80    // the incremental path.
81    let mut full_ids = ids.clone();
82    let mut full_codes = Vec::new();
83    for _ in 0..20 {
84        let ids_i64: Vec<i64> = full_ids.iter().map(|&id| i64::from(id)).collect();
85        let len = ids_i64.len();
86        let input = Tensor::<B, 2, Int>::from_data(TensorData::new(ids_i64, [1, len]), &device);
87        let hidden = model.forward(input, true);
88        let last = hidden.narrow(1, len - 1, 1).reshape([1, hidden_dim]);
89        let logits = last
90            .matmul(weight.clone().transpose())
91            .reshape([config.vocab_size as usize]);
92        let values: Vec<f32> = logits
93            .into_data()
94            .convert::<f32>()
95            .to_vec()
96            .map_err(|e| anyhow::anyhow!("{e}"))?;
97        let next = values
98            .iter()
99            .enumerate()
100            .max_by(|(_, a), (_, b)| a.total_cmp(b))
101            .map(|(i, _)| i as u32)
102            .unwrap();
103        if next == AudioCodeVocab::IM_END_ID {
104            break;
105        }
106        if let Some(code) = vocab.token_id_to_code(next) {
107            full_codes.push(code);
108        }
109        full_ids.push(next);
110    }
111    println!(
112        "\nfull-forward greedy: {} codes: {full_codes:?}",
113        full_codes.len()
114    );
115    Ok(())
116}
lm_forward_probe/lm_forward_probe.rs

lm_forward_probe/
lm_forward_probe.rs
