metrale_model_layers/layers/ngram_embed/
ids.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The n-gram id core: the EOS-aware right shift and the
4//! polynomial rolling hash, on token ids only (no device, no weights).
5//!
6//! Owner: model-layers (n-gram embedding).
7//! Invariants: none beyond the types.
8
9use super::NgramDims;
10
11/// 2026-09-25: `out[t] = ctx[t - n]` when `t - n` lies in the same segment
12/// as `t` (a segment ends at an EOS token, inclusive), else 0. A segment of
13/// n tokens or fewer is all zeros.
14fn shift_right_ignore_eos(ctx: &[u32], n: usize, eos: u32) -> Vec<u32> {
15    let len = ctx.len();
16    let mut out = vec![0u32; len];
17    let mut prev = 0usize;
18    for (pos, &tok) in ctx.iter().enumerate() {
19        if tok == eos {
20            let end = pos + 1;
21            if end - prev > n {
22                out[prev + n..end].copy_from_slice(&ctx[prev..end - n]);
23            }
24            prev = end;
25        }
26    }
27    if prev < len && len - prev > n {
28        out[prev + n..len].copy_from_slice(&ctx[prev..len - n]);
29    }
30    out
31}
32
33/// 2026-09-25: Row ids of every table over `ctx`. Returns `num_tables()`
34/// vectors of `ctx.len()` ids each, in table index order
35/// (`(ngram-2)*K + split`).
36pub fn ngram_ids(dims: &NgramDims, ctx: &[u32]) -> Vec<Vec<u64>> {
37    let mut out = Vec::with_capacity(dims.num_tables());
38    // 2026-09-25: One shift per d in 1..N, shared by every table.
39    let shifts: Vec<Vec<u32>> = (1..dims.neighbor_num)
40        .map(|d| shift_right_ignore_eos(ctx, d, dims.eos_token_id))
41        .collect();
42    for ngram in 2..=dims.neighbor_num {
43        for split in 0..dims.split_num {
44            let index = (ngram - 2) * dims.split_num + split;
45            let t = dims.table_rows(index);
46            let mods = dims.vocab_mods(ngram, split);
47            let ids = ctx
48                .iter()
49                .enumerate()
50                .map(|(pos, &x)| {
51                    let mut acc = x as u64;
52                    for (d, &m) in mods.iter().enumerate() {
53                        acc += shifts[d][pos] as u64 * m;
54                    }
55                    acc % t
56                })
57                .collect();
58            out.push(ids);
59        }
60    }
61    out
62}
metrale_model_layers/layers/ngram_embed/ids.rs

metrale_model_layers/layers/ngram_embed/
ids.rs
