geographdb-core 0.5.4

//! Per-position hidden-state divergence for the approximate-transformer
//! tiebreak experiment.
//!
//! For each requested context length and geometric-only fraction, this example
//! compares the hidden state at *every* position against the full hybrid
//! baseline.  The output CSV can be visualised with
//! `examples/plot_per_position_divergence.py`.
//!
//! Optional arguments:
//!   --checkpoint <path>        load trained weights from a `flatten_params` dump
//!   --lengths <64,128,256>     comma-separated context lengths to evaluate
//!   --fractions <0,0.5,1>      comma-separated geometric-only fractions
//!   --seed <42>                RNG seed for the model init (ignored if checkpoint)
//!   --mode <random|last-hybrid> mask mode: random positions, or last token always hybrid
//!   --neighbor-mode <ring|random> how to choose geometric neighbours per token
//!   --vocab-size <256>         model vocabulary size
//!   --embed-dim <64>           attention embedding dimension
//!   --hidden-dim <128>         MLP hidden dimension
//!   --output-dim <64>          MLP output dimension
//!   --num-neighbors <4>        number of geometric neighbours per token
//!   --plasticity <false>       enable learnable edge weights

use geographdb_core::algorithms::attention::GraphAttentionClassifier;
use geographdb_core::algorithms::attention_divergence::measure_divergence_per_position;

#[derive(Debug, Clone, Copy)]
enum MaskMode {
    Random,
    LastHybrid,
}

#[derive(Debug, Clone, Copy)]
enum NeighborMode {
    Ring,
    Random,
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let checkpoint = arg_value(&args, "--checkpoint");
    let lengths: Vec<usize> = arg_value(&args, "--lengths")
        .unwrap_or_else(|| "64".to_string())
        .split(',')
        .map(|s| {
            s.trim()
                .parse()
                .expect("--lengths must be a comma-separated list of positive integers")
        })
        .collect();
    let fractions: Vec<f32> = arg_value(&args, "--fractions")
        .unwrap_or_else(|| "0.0,0.25,0.5,0.75,1.0".to_string())
        .split(',')
        .map(|s| {
            s.trim()
                .parse()
                .expect("--fractions must be a comma-separated list of floats")
        })
        .collect();
    let seed = arg_value(&args, "--seed")
        .and_then(|s| s.parse().ok())
        .unwrap_or(42u32);
    let mode = match arg_value(&args, "--mode")
        .unwrap_or_else(|| "last-hybrid".to_string())
        .to_lowercase()
        .as_str()
    {
        "last-hybrid" => MaskMode::LastHybrid,
        "random" => MaskMode::Random,
        other => panic!("unknown --mode: {other}. Use 'random' or 'last-hybrid'"),
    };
    let neighbor_mode = match arg_value(&args, "--neighbor-mode")
        .unwrap_or_else(|| "ring".to_string())
        .to_lowercase()
        .as_str()
    {
        "ring" => NeighborMode::Ring,
        "random" => NeighborMode::Random,
        other => panic!("unknown --neighbor-mode: {other}. Use 'ring' or 'random'"),
    };

    let vocab_size = arg_value(&args, "--vocab-size")
        .and_then(|s| s.parse().ok())
        .unwrap_or(256usize);
    let embed_dim = arg_value(&args, "--embed-dim")
        .and_then(|s| s.parse().ok())
        .unwrap_or(64usize);
    let hidden_dim = arg_value(&args, "--hidden-dim")
        .and_then(|s| s.parse().ok())
        .unwrap_or(128usize);
    let output_dim = arg_value(&args, "--output-dim")
        .and_then(|s| s.parse().ok())
        .unwrap_or(64usize);
    let num_neighbors = arg_value(&args, "--num-neighbors")
        .and_then(|s| s.parse().ok())
        .unwrap_or(4usize);
    let plasticity = arg_value(&args, "--plasticity")
        .map(|s| matches!(s.to_lowercase().as_str(), "true" | "1" | "yes"))
        .unwrap_or(false);

    let mut model = GraphAttentionClassifier::new(
        vocab_size,
        embed_dim,
        hidden_dim,
        output_dim,
        num_neighbors,
        seed,
        Some(10000.0),
        plasticity,
    );

    if let Some(path) = checkpoint {
        let bytes = std::fs::read(&path).expect("failed to read checkpoint");
        let params: Vec<f32> = bytes
            .chunks_exact(4)
            .map(|b| f32::from_le_bytes(b.try_into().expect("checkpoint size not a multiple of 4")))
            .collect();
        model.load_flat_params(&params);
        eprintln!(
            "Loaded checkpoint: {} ({} params)  arch={}x{}x{}x{} k={} plasticity={}",
            path,
            params.len(),
            vocab_size,
            embed_dim,
            hidden_dim,
            output_dim,
            num_neighbors,
            plasticity
        );
    }

    let mode_str = match mode {
        MaskMode::Random => "random",
        MaskMode::LastHybrid => "last_hybrid",
    };
    let neighbor_mode_str = match neighbor_mode {
        NeighborMode::Ring => "ring",
        NeighborMode::Random => "random",
    };
    println!("neighbor_mode,mode,length,geometric_fraction,position,geo_l2,geo_mse,geo_mae,geo_cosine,mixed_l2,mixed_mse,mixed_mae,mixed_cosine");

    for &n_context in &lengths {
        let token_ids: Vec<u32> = (0..n_context).map(|i| (i % vocab_size) as u32).collect();
        let neighbors_owned: Vec<Vec<usize>> = match neighbor_mode {
            NeighborMode::Ring => token_ids
                .iter()
                .map(|&t| {
                    (0..num_neighbors)
                        .map(|k| (t as usize + k + 1) % vocab_size)
                        .collect()
                })
                .collect(),
            NeighborMode::Random => {
                let mut state = 12345u32;
                token_ids
                    .iter()
                    .map(|_| {
                        let mut placed = 0usize;
                        let mut nbrs = Vec::with_capacity(num_neighbors);
                        while placed < num_neighbors {
                            state ^= state << 13;
                            state ^= state >> 17;
                            state ^= state << 5;
                            let idx = (state as usize) % vocab_size;
                            if !nbrs.contains(&idx) {
                                nbrs.push(idx);
                                placed += 1;
                            }
                        }
                        nbrs
                    })
                    .collect()
            }
        };
        let neighbors: Vec<&[usize]> = neighbors_owned.iter().map(|v| v.as_slice()).collect();

        for &fraction in &fractions {
            let mask = make_mask(n_context, fraction, 12345, mode);
            let div = measure_divergence_per_position(&model, &token_ids, &neighbors, &mask);

            for (geo, mixed) in div
                .geometric_vs_hybrid
                .iter()
                .zip(div.mixed_vs_hybrid.iter())
            {
                println!(
                    "{},{},{},{:.2},{},{:.10e},{:.10e},{:.10e},{:.10e},{:.10e},{:.10e},{:.10e},{:.10e}",
                    neighbor_mode_str,
                    mode_str,
                    n_context,
                    fraction,
                    geo.position,
                    geo.hidden_l2,
                    geo.hidden_mse,
                    geo.hidden_mae,
                    geo.cosine_similarity,
                    mixed.hidden_l2,
                    mixed.hidden_mse,
                    mixed.hidden_mae,
                    mixed.cosine_similarity,
                );
            }
        }
    }
}

fn make_mask(n: usize, fraction_true: f32, seed: u32, mode: MaskMode) -> Vec<bool> {
    match mode {
        MaskMode::Random => random_mask(n, fraction_true, seed),
        MaskMode::LastHybrid => last_hybrid_mask(n, fraction_true),
    }
}

fn last_hybrid_mask(n: usize, fraction_true: f32) -> Vec<bool> {
    let mut mask = vec![false; n];
    if n == 0 {
        return mask;
    }
    let n_true = ((fraction_true * n as f32).round() as usize).min(n.saturating_sub(1));
    for item in mask.iter_mut().take(n_true) {
        *item = true;
    }
    mask[n - 1] = false;
    mask
}

fn random_mask(n: usize, fraction_true: f32, seed: u32) -> Vec<bool> {
    let mut mask = vec![false; n];
    let n_true = (fraction_true * n as f32).round() as usize;
    if n_true == 0 {
        return mask;
    }
    if n_true >= n {
        return vec![true; n];
    }

    let mut state = seed;
    let mut placed = 0;
    while placed < n_true {
        state ^= state << 13;
        state ^= state >> 17;
        state ^= state << 5;
        let idx = (state as usize) % n;
        if !mask[idx] {
            mask[idx] = true;
            placed += 1;
        }
    }
    mask
}

fn arg_value(args: &[String], key: &str) -> Option<String> {
    args.windows(2).find(|w| w[0] == key).map(|w| w[1].clone())
}