memra-engine 0.125.0

//! FR-Spec ranking builder: tokenize a corpus with a model's OWN tokenizer, rank token ids by
//! frequency, and emit a minimal GGUF carrying the top-N `d2t` (i32) list that MEMRA_FRSPEC_TRIM /
//! MEMRA_MTP_DRAFT consume. Needed because trim files are VOCAB artifacts — the published Qwen3.6
//! rankings from another tokenizer cannot transfer to Hy3's vocabulary.
//!
//! usage: frspec-rank <model.gguf|hf_dir> <out.gguf> <topN>
//!                    [--coverage-ranks ranks.txt] <corpus file/dir>...
//!
//! Accepts EITHER a .gguf file (tokenizer from GGUF metadata) OR an HF safetensors directory
//! containing tokenizer.json (the ST-native path — no GGUF dependency for ST checkpoints).
use memra_gguf::GgufFile;
use memra_tokenizer::Tokenizer;

fn collect_files(path: &std::path::Path, out: &mut Vec<std::path::PathBuf>) {
    let Ok(meta) = std::fs::symlink_metadata(path) else {
        return;
    };
    // Corpus roots often contain convenience symlinks back into a projects tree. Following those
    // makes the walk cyclic and silently weights the same source multiple times.
    if meta.file_type().is_symlink() {
        return;
    }
    if path.is_dir() {
        if matches!(
            path.file_name().and_then(|name| name.to_str()),
            Some(".git" | ".venv" | "node_modules" | "target")
        ) {
            return;
        }
        if let Ok(rd) = std::fs::read_dir(path) {
            for e in rd.flatten() {
                collect_files(&e.path(), out);
            }
        }
    } else if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
        // text-ish sources only
        if matches!(
            ext,
            "txt"
                | "md"
                | "rs"
                | "py"
                | "cu"
                | "c"
                | "h"
                | "cpp"
                | "json"
                | "toml"
                | "sh"
                | "js"
                | "ts"
        ) {
            out.push(path.to_path_buf());
        }
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 5 {
        eprintln!(
            "usage: frspec-rank <model.gguf|hf_dir> <out.gguf> <topN> \
             [--coverage-ranks ranks.txt] <corpus>..."
        );
        std::process::exit(1);
    }
    // DIRECTORY path = HF safetensors checkpoint (tokenizer.json); file = GGUF.
    let model_path = std::path::Path::new(&args[1]);
    let is_dir = model_path.is_dir();
    let g: Option<GgufFile> = if is_dir {
        None
    } else {
        Some(GgufFile::open(&args[1])?)
    };
    let tok = match &g {
        Some(g) => Tokenizer::from_gguf(g).map_err(|e| format!("tokenizer: {e}"))?,
        None => Tokenizer::from_hf_dir(model_path)
            .map_err(|e| format!("HF tokenizer init failed: {e}"))?,
    };
    let top_n: usize = args[3].parse()?;
    let vocab = tok.vocab_size();
    let mut coverage_ranks: Option<&str> = None;
    let mut corpus_roots: Vec<&str> = Vec::new();
    let mut index = 4;
    while index < args.len() {
        if args[index] == "--coverage-ranks" {
            let path = args
                .get(index + 1)
                .ok_or("--coverage-ranks requires a path")?;
            coverage_ranks = Some(path);
            index += 2;
        } else {
            corpus_roots.push(&args[index]);
            index += 1;
        }
    }
    if corpus_roots.is_empty() {
        return Err("at least one corpus file or directory is required".into());
    }

    let mut files = Vec::new();
    for a in corpus_roots {
        collect_files(std::path::Path::new(a), &mut files);
    }
    eprintln!(
        "[frspec-rank] {} corpus files, vocab {}",
        files.len(),
        vocab
    );

    let mut counts = vec![0u64; vocab];
    let mut total = 0u64;
    for f in &files {
        let Ok(text) = std::fs::read_to_string(f) else {
            continue;
        };
        for id in tok.encode(&text, false) {
            if (id as usize) < vocab {
                counts[id as usize] += 1;
                total += 1;
            }
        }
    }
    let distinct = counts.iter().filter(|&&count| count > 0).count();
    eprintln!("[frspec-rank] {total} tokens counted ({distinct} distinct ids)");

    // rank by frequency desc, id asc tiebreak (deterministic). Zero-count ids are EXCLUDED from
    // preference but the list must still fill top_n — pad with ascending unseen ids (they cost
    // nothing: the draft never proposes what the head's trimmed rows can't produce... they simply
    // occupy cover slots). Practical corpora cover ~60-120k distinct ids.
    let d2t = memra_gguf::d2t::rank_top_n(&counts, top_n);
    let covered: u64 = d2t.iter().map(|&i| counts[i as usize]).sum();
    eprintln!(
        "[frspec-rank] top {} covers {:.2}% of corpus tokens",
        d2t.len(),
        covered as f64 / total.max(1) as f64 * 100.0
    );
    if let Some(path) = coverage_ranks {
        let mut seen = vec![false; vocab];
        let mut supplied = Vec::new();
        for (line_number, line) in std::fs::read_to_string(path)?.lines().enumerate() {
            if line.trim().is_empty() {
                continue;
            }
            let token: usize = line
                .trim()
                .parse()
                .map_err(|error| format!("{path}:{}: {error}", line_number + 1))?;
            if token >= vocab {
                return Err(
                    format!("{path}:{}: token {token} >= vocab {vocab}", line_number + 1).into(),
                );
            }
            if seen[token] {
                return Err(format!("{path}:{}: duplicate token {token}", line_number + 1).into());
            }
            seen[token] = true;
            supplied.push(token);
        }
        let supplied_covered: u64 = supplied.iter().map(|&token| counts[token]).sum();
        eprintln!(
            "[frspec-rank] supplied map {} ({} ids) covers {:.2}% of corpus tokens",
            path,
            supplied.len(),
            supplied_covered as f64 / total.max(1) as f64 * 100.0
        );
        let mut missing: Vec<(usize, u64)> = counts
            .iter()
            .copied()
            .enumerate()
            .filter(|&(token, count)| count > 0 && !seen[token])
            .collect();
        missing.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
        for (rank, (token, count)) in missing.into_iter().take(20).enumerate() {
            eprintln!(
                "[frspec-rank] uncovered#{:02} id={} count={} piece={:?}",
                rank + 1,
                token,
                count,
                tok.decode(&[token as u32])
            );
        }
    }
    memra_gguf::d2t::write_d2t(&args[2], &d2t)?;
    eprintln!(
        "[frspec-rank] wrote {} ({} ids) + {}.txt",
        args[2],
        d2t.len(),
        args[2]
    );
    Ok(())
}