searchi/tools/speech-rank/src/alignment.rs

//! Independent reconstruction of the frozen ordinary token-time experiment.
use super::{number, parsed_segments};
use anyhow::{ensure, Context, Result};
use serde_json::{json, Value};

pub const POLICY: &str =
    "ordinary-token-sentences-v1;punctuation-boundaries;exact-text;no-dtw;no-labels";
pub const POLICY_SHA: &str = "5bc603e792a3abcc2370d4d1f7a13f538f808fb3b959b506a16d5823d3d7a126";
pub const CONTRACT_SHA: &str = "aac049c5bf3e17b07f4f2f7d92bce18c27aef5a67364181e47815301e75ca520";
pub const DTW_POLICY: &str =
    "dtw-token-sentences-v1;centisecond-points;first-last-point-span;exact-text;no-labels";
pub const DTW_POLICY_SHA: &str = "0ea19b1a2ee9da75f0284de33aeca0fcdff0a77d88d232b0b25821a7115b83b5";
pub const DTW_CONTRACT_SHA: &str =
    "61594d0e77c561d5d3b7c9e63178a8123841516b8f8fa781458b21ab717ad47f";

pub fn enabled(asr: &Value) -> Result<bool> {
    if asr.get("alignment_policy").is_none() {
        ensure!(
            asr.get("alignment_contract_sha256").is_none()
                && asr.get("alignment_policy_sha256").is_none(),
            "Partial alignment provenance"
        );
        return Ok(false);
    }
    ensure!(
        ((asr["alignment_policy"] == POLICY
            && asr["alignment_policy_sha256"] == POLICY_SHA
            && asr["alignment_contract_sha256"] == CONTRACT_SHA)
            || (asr["alignment_policy"] == DTW_POLICY
                && asr["alignment_policy_sha256"] == DTW_POLICY_SHA
                && asr["alignment_contract_sha256"] == DTW_CONTRACT_SHA))
            && asr["candidate"] == "base"
            && asr["worker_sha256"]
                == "12c681b4bcd9c7679b647b9a041e121a2f9e72cb2a9c6f189f529ec8b79c36f8"
            && asr["policy"] == super::POLICY,
        "Unknown token alignment provenance"
    );
    Ok(true)
}

fn boundary(text: &str) -> bool {
    // Python's frozen regex `$` also matches before a final newline.
    let text = text.strip_suffix('\n').unwrap_or(text);
    text.trim_end_matches(['"', '\u{201d}', '\u{2019}', '\'', ')', ']'])
        .ends_with(['.', '!', '?', '。', '！', '？'])
}

pub fn reconstruct(raw: &Value, offset: u64, duration: u64) -> Result<Vec<Value>> {
    reconstruct_mode(raw, offset, duration, false)
}

pub fn reconstruct_mode(raw: &Value, offset: u64, duration: u64, dtw: bool) -> Result<Vec<Value>> {
    let originals = parsed_segments(raw, offset, duration)?;
    let rows = raw["transcription"]
        .as_array()
        .context("Missing transcription")?;
    let mut result = Vec::new();
    let mut original_index = 0;
    let mut dtw_previous = 0;
    for row in rows {
        let text = row["text"].as_str().context("Missing original text")?;
        if text.trim().is_empty() {
            continue;
        }
        let original = &originals[original_index];
        let tokens = row["tokens"].as_array().context("Missing token times")?;
        ensure!(
            !tokens.is_empty() && tokens.len() <= 4096,
            "Unbounded token times"
        );
        let mut exact = String::new();
        let mut phrase = String::new();
        let mut phrase_start = None;
        let mut phrase_end = 0;
        let mut previous = 0;
        for token in tokens {
            let id = number(&token["id"])?;
            let token_text = token["text"].as_str().context("Invalid token text")?;
            if id >= 50257 {
                ensure!(
                    token_text.starts_with("[_") && token_text.ends_with(']'),
                    "Invalid special token"
                );
                continue;
            }
            let (start, end) = if dtw {
                let point = number(&token["t_dtw"])?
                    .checked_mul(10)
                    .context("DTW point overflow")?;
                ensure!(
                    point >= dtw_previous && point <= duration,
                    "Invalid DTW source point"
                );
                dtw_previous = point;
                (point, point)
            } else {
                (
                    number(&token["offsets"]["from"])?,
                    number(&token["offsets"]["to"])?,
                )
            };
            ensure!(
                start >= previous
                    && start <= end
                    && end <= duration
                    && (dtw
                        || (start >= number(&row["offsets"]["from"])?
                            && end <= number(&row["offsets"]["to"])?)),
                "Invalid ordinary token time"
            );
            previous = start;
            exact.push_str(token_text);
            ensure!(exact.len() <= 4096, "Oversized token text");
            phrase_start.get_or_insert(start);
            phrase_end = end;
            phrase.push_str(token_text);
            if boundary(token_text) {
                append(
                    &mut result,
                    &phrase,
                    phrase_start,
                    phrase_end,
                    offset,
                    original_index,
                    original,
                )?;
                phrase.clear();
                phrase_start = None;
            }
        }
        ensure!(
            exact == text,
            "Token text differs from original recognized segment"
        );
        append(
            &mut result,
            &phrase,
            phrase_start,
            phrase_end,
            offset,
            original_index,
            original,
        )?;
        original_index += 1;
    }
    ensure!(result.len() <= 10000, "Oversized aligned passages");
    Ok(result)
}

fn append(
    result: &mut Vec<Value>,
    text: &str,
    start: Option<u64>,
    end: u64,
    offset: u64,
    index: usize,
    original: &Value,
) -> Result<()> {
    if !text.trim().is_empty() {
        result.push(json!({"start_ms":offset.checked_add(start.context("Missing phrase start")?).context("Timestamp overflow")?,
            "end_ms":offset.checked_add(end).context("Timestamp overflow")?, "text":text,
            "original_segment_index":index,"original_start_ms":original["start_ms"],"original_end_ms":original["end_ms"]}));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    fn raw() -> Value {
        json!({"transcription":[{"text":" Hello. Guten Tag!", "offsets":{"from":0,"to":2000},
          "tokens":[{"id":50364,"text":"[_BEG_]"},
            {"id":100,"text":" Hello.","offsets":{"from":100,"to":900}},
            {"id":101,"text":" Guten","offsets":{"from":1100,"to":1500}},
            {"id":102,"text":" Tag!","offsets":{"from":1500,"to":1900}}]}]})
    }
    #[test]
    fn reconstructs_exact_sentences_and_absolute_source_coordinates() {
        let result = reconstruct(&raw(), 30000, 2000).unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(
            result[0],
            json!({"start_ms":30100,"end_ms":30900,"text":" Hello.","original_segment_index":0,"original_start_ms":30000,"original_end_ms":32000})
        );
        assert_eq!(result[1]["text"], " Guten Tag!");
        assert_eq!(result[1]["start_ms"], 31100);
        assert!(reconstruct(&raw(), u64::MAX, 2000).is_err());
    }
    #[test]
    fn rejects_missing_times_outside_bounds_nonmonotonic_and_changed_text() {
        for replacement in [
            json!({"id":100,"text":" Hello."}),
            json!({"id":100,"text":" Hello.","offsets":{"from":0,"to":3000}}),
            json!({"id":100,"text":" changed","offsets":{"from":100,"to":900}}),
            json!({"id":50364,"text":"real words"}),
        ] {
            let mut changed = raw();
            changed["transcription"][0]["tokens"][1] = replacement;
            assert!(reconstruct(&changed, 0, 2000).is_err());
        }
        let mut changed = raw();
        changed["transcription"][0]["tokens"][3]["offsets"]["from"] = json!(1000);
        assert!(reconstruct(&changed, 0, 2000).is_err());
    }
    #[test]
    fn provenance_cannot_be_partial_unknown_or_small_model() {
        assert!(!enabled(&json!({})).unwrap());
        assert!(enabled(&json!({"alignment_contract_sha256":CONTRACT_SHA})).is_err());
        let report = json!({"alignment_policy":POLICY,"alignment_policy_sha256":POLICY_SHA,
            "alignment_contract_sha256":CONTRACT_SHA,"candidate":"base",
            "worker_sha256":"12c681b4bcd9c7679b647b9a041e121a2f9e72cb2a9c6f189f529ec8b79c36f8",
            "policy":super::super::POLICY});
        assert!(enabled(&report).unwrap());
        let mut changed = report;
        changed["candidate"] = json!("small");
        assert!(enabled(&changed).is_err());
    }
    #[test]
    fn dtw_points_are_independent_bounded_monotonic_centiseconds() {
        let mut raw = raw();
        for (index, point) in [(1, 20), (2, 70), (3, 90)] {
            raw["transcription"][0]["tokens"][index]["t_dtw"] = json!(point);
        }
        raw["transcription"][0]["tokens"][1]["offsets"] = json!({"from":900,"to":100});
        assert!(reconstruct(&raw, 0, 2000).is_err());
        let result = reconstruct_mode(&raw, 30000, 2000, true).unwrap();
        assert_eq!(result[0]["start_ms"], 30200);
        assert_eq!(result[0]["end_ms"], 30200);
        assert_eq!(result[1]["start_ms"], 30700);
        assert_eq!(result[1]["end_ms"], 30900);
        for point in [json!(-1), json!(201), json!(1.5), json!(10), Value::Null] {
            let mut changed = raw.clone();
            changed["transcription"][0]["tokens"][2]["t_dtw"] = point;
            assert!(reconstruct_mode(&changed, 0, 2000, true).is_err());
        }
    }
}