Validate finer speech passages independently before Google ranking

tools/speech-rank/src/alignment.rs

@@ -0,0 +1,190 @@
+//! 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 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["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>> {
+    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;
+    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 = number(&token["offsets"]["from"])?;
+            let end = number(&token["offsets"]["to"])?;
+            ensure!(
+                start >= previous
+                    && start <= end
+                    && end <= duration
+                    && 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());
+    }
+}