238 lines
9.3 KiB
Rust
238 lines
9.3 KiB
Rust
//! 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());
|
||
}
|
||
}
|
||
}
|