fxtranslate 0.4.2

#![cfg(not(target_arch = "wasm32"))] // native-only (filesystem fixtures / oracle); wasm runs gemm_parity
//! Batched decoder parity (batch-invariance).
//!
//! `greedy_batch` / `translate_batch` must produce, for each sentence, exactly
//! what the single-sentence `greedy` / `translate` produces — decoder rows are
//! independent (SSRU cell per row, cross-attention to each sentence's own masked
//! context, no decoder self-attention), so batching a block must not change any
//! sentence's tokens. The single path is validated against the marian reference
//! trace, so this is cheat-proof: a cross-attention-mask or per-row EOS bug shows
//! up as a per-sentence divergence, not a comparison against our own batch output.
//!
//! Skips (rather than fails) when the en-fr model isn't downloaded.

use fxtranslate::engine::Engine;

const MODEL: &str = "../../../data/models/enfr/model.enfr.intgemm.alphas.bin";
const VOCAB: &str = "../../../data/models/enfr/vocab.enfr.spm";

fn engine() -> Option<Engine> {
    if !std::path::Path::new(MODEL).exists() {
        eprintln!("skipping batched-decode parity: {MODEL} absent");
        return None;
    }
    Some(Engine::load(MODEL, VOCAB, VOCAB).expect("engine loads"))
}

#[test]
fn greedy_batch_matches_single_per_sentence() {
    let Some(eng) = engine() else { return };
    // Mixed lengths so the block pads (source) and sentences finish at different
    // steps (target) — exercises the cross-attention mask and per-row EOS.
    let texts = [
        "The cat sat on the mat.",
        "Dogs run.",
        "Scientists carefully explained the experiment to the students.",
        "Birds fly south.",
    ];
    let ids: Vec<Vec<u32>> = texts.iter().map(|t| eng.src_ids(t)).collect();

    let batched = eng.greedy_batch(&ids);
    for (b, sid) in ids.iter().enumerate() {
        let single = eng.greedy(sid);
        eprintln!(
            "sentence {b}: single {} toks, batched {} toks",
            single.len(),
            batched[b].len()
        );
        assert_eq!(
            batched[b], single,
            "sentence {b}: batched greedy diverges from single-sentence greedy"
        );
    }
}

#[test]
fn translate_batch_matches_single() {
    let Some(eng) = engine() else { return };
    let texts = [
        "Hello world.",
        "The quick brown fox jumps.",
        "Good morning.",
    ];
    let batched = eng.translate_batch(&texts);
    for (b, t) in texts.iter().enumerate() {
        assert_eq!(
            batched[b],
            eng.translate(t),
            "sentence {b}: translate_batch diverges from translate"
        );
    }
}

/// The data-parallel batch path (feature `threads`) partitions the batch across
/// worker threads sharing one copy of the weights. Every sentence must decode to
/// exactly what the serial batch produces — the split is over independent
/// sentences, so threading must not move a single token. More sentences than
/// workers so the partition spans multiple chunks per thread.
#[cfg(feature = "threads")]
#[test]
fn greedy_batch_parallel_matches_serial() {
    let Some(eng) = engine() else { return };
    let texts = [
        "The cat sat on the mat.",
        "Dogs run.",
        "Scientists carefully explained the experiment to the students.",
        "Birds fly south.",
        "Hello world.",
        "The quick brown fox jumps over the lazy dog.",
        "Good morning.",
        "She sells seashells by the seashore.",
        "Winter is coming soon.",
        "They travelled across the country by train.",
    ];
    let ids: Vec<Vec<u32>> = texts.iter().map(|t| eng.src_ids(t)).collect();
    let serial = eng.greedy_batch(&ids); // default threads == 1
    let eng = eng.with_threads(8);
    let parallel = eng.greedy_batch(&ids);
    assert_eq!(
        parallel, serial,
        "data-parallel greedy_batch diverges from the serial batch"
    );
}