cuepack: cues into passages that keep their timestamps

crates/cuepack/src/lib.rs

@@ -0,0 +1,197 @@
+//! Packs cues into passages worth embedding.
+//!
+//! A cue is a few seconds of speech and half a sentence. An embedding of one
+//! says almost nothing, and a search over them returns fragments. Packing
+//! them into a minute or so of continuous speech gives a passage with enough
+//! context to mean something, and — because the pack keeps the first cue's
+//! start and the last cue's end — one that can still be cited to the second.
+
+use corpus_core::Chunk;
+use cuekit::Cue;
+
+/// Where a passage is flushed, whichever comes first.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct Budgets {
+    /// Seconds of audio a passage covers.
+    ///
+    /// Seventy-five is about a paragraph of speech: long enough to hold an
+    /// argument, short enough that a result points at a place in the
+    /// recording rather than at a region of it.
+    pub seconds: f64,
+    /// Tokens a passage may hold, whatever the clock says.
+    pub max_tokens: usize,
+}
+
+impl Default for Budgets {
+    fn default() -> Self {
+        Self {
+            seconds: 75.0,
+            max_tokens: 512,
+        }
+    }
+}
+
+/// Pack consecutive cues into passages.
+///
+/// A passage is flushed when the window from its first cue's start to the
+/// current cue's end reaches the seconds budget, or when adding the current
+/// cue would take it past the token budget.
+///
+/// Immediately repeated lines are dropped rather than packed, because
+/// automatic captioning repeats each line across overlapping windows and a
+/// passage padded with its own echo embeds as something nobody said.
+pub fn pack(cues: &[Cue], budgets: &Budgets) -> Vec<Chunk> {
+    let mut passages: Vec<Chunk> = Vec::new();
+    let mut buffered: Vec<&Cue> = Vec::new();
+    let mut buffered_tokens = 0usize;
+    let mut last_text: Option<&str> = None;
+
+    for cue in cues {
+        let text = cue.text.trim_matches(textkit::is_space);
+        if text.is_empty() || Some(text) == last_text {
+            continue;
+        }
+        last_text = Some(text);
+
+        let tokens = textkit::count(text);
+        let span = buffered
+            .first()
+            .map_or(0.0, |first: &&Cue| cue.end - first.start);
+        if !buffered.is_empty()
+            && (span >= budgets.seconds || buffered_tokens + tokens > budgets.max_tokens)
+        {
+            flush(&mut passages, &mut buffered, &mut buffered_tokens);
+        }
+        buffered.push(cue);
+        buffered_tokens += tokens;
+    }
+    flush(&mut passages, &mut buffered, &mut buffered_tokens);
+    passages
+}
+
+fn flush(passages: &mut Vec<Chunk>, buffered: &mut Vec<&Cue>, buffered_tokens: &mut usize) {
+    if buffered.is_empty() {
+        return;
+    }
+    let text = buffered
+        .iter()
+        .map(|cue| cue.text.trim_matches(textkit::is_space))
+        .filter(|text| !text.is_empty())
+        .collect::<Vec<_>>()
+        .join(" ");
+
+    if !text.is_empty() {
+        let tokens = textkit::count(&text);
+        let mut chunk = Chunk::new(passages.len() as u32, text, tokens as u32);
+        chunk.start_time = buffered.first().map(|cue| cue.start);
+        chunk.end_time = buffered.last().map(|cue| cue.end);
+        passages.push(chunk);
+    }
+    buffered.clear();
+    *buffered_tokens = 0;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn cue(start: f64, end: f64, text: &str) -> Cue {
+        Cue::new(start, end, text)
+    }
+
+    #[test]
+    fn cues_within_the_budgets_make_one_passage() {
+        let cues = [
+            cue(0.0, 5.0, "the mill stood"),
+            cue(5.0, 10.0, "at the northern end"),
+        ];
+        let packed = pack(&cues, &Budgets::default());
+        assert_eq!(packed.len(), 1);
+        assert_eq!(packed[0].text, "the mill stood at the northern end");
+        assert_eq!(packed[0].start_time, Some(0.0));
+        assert_eq!(packed[0].end_time, Some(10.0));
+        assert_eq!(packed[0].idx, 0);
+    }
+
+    #[test]
+    fn a_passage_is_flushed_when_it_has_run_long_enough() {
+        // The window is measured to the end of the cue being considered, so
+        // the cue that trips the budget starts the next passage rather than
+        // finishing this one.
+        let cues = [
+            cue(0.0, 40.0, "one"),
+            cue(40.0, 80.0, "two"),
+            cue(80.0, 90.0, "three"),
+        ];
+        let packed = pack(&cues, &Budgets::default());
+        assert_eq!(packed.len(), 2);
+        assert_eq!(packed[0].text, "one");
+        assert_eq!(packed[0].end_time, Some(40.0));
+        assert_eq!(packed[1].text, "two three");
+        assert_eq!(packed[1].idx, 1);
+    }
+
+    #[test]
+    fn a_passage_is_flushed_before_it_overruns_the_token_budget() {
+        let budgets = Budgets {
+            seconds: 1000.0,
+            max_tokens: 6,
+        };
+        let cues = [
+            cue(0.0, 1.0, "one two three"),
+            cue(1.0, 2.0, "four five six"),
+            cue(2.0, 3.0, "seven eight nine"),
+        ];
+        let packed = pack(&cues, &budgets);
+        assert_eq!(packed.len(), 2);
+        for passage in &packed {
+            assert!(passage.tokens <= 6, "{} tokens", passage.tokens);
+        }
+    }
+
+    #[test]
+    fn an_immediately_repeated_line_is_dropped() {
+        let cues = [
+            cue(0.0, 2.0, "the mill stood"),
+            cue(1.0, 3.0, "the mill stood"),
+            cue(3.0, 4.0, "at the northern end"),
+        ];
+        let packed = pack(&cues, &Budgets::default());
+        assert_eq!(packed[0].text, "the mill stood at the northern end");
+    }
+
+    #[test]
+    fn a_line_that_comes_back_later_is_kept() {
+        let cues = [
+            cue(0.0, 1.0, "yes"),
+            cue(1.0, 2.0, "no"),
+            cue(2.0, 3.0, "yes"),
+        ];
+        assert_eq!(pack(&cues, &Budgets::default())[0].text, "yes no yes");
+    }
+
+    #[test]
+    fn an_empty_cue_is_not_a_break_in_the_repetition() {
+        let cues = [
+            cue(0.0, 1.0, "same"),
+            cue(1.0, 2.0, "   "),
+            cue(2.0, 3.0, "same"),
+        ];
+        assert_eq!(pack(&cues, &Budgets::default())[0].text, "same");
+    }
+
+    #[test]
+    fn a_passage_carries_times_that_can_be_cited() {
+        let cues = [cue(12.5, 15.25, "the mill stood")];
+        let packed = pack(&cues, &Budgets::default());
+        assert_eq!(packed[0].start_time, Some(12.5));
+        assert_eq!(packed[0].end_time, Some(15.25));
+        assert_eq!(packed[0].page, None);
+    }
+
+    #[test]
+    fn nothing_in_nothing_out() {
+        assert!(pack(&[], &Budgets::default()).is_empty());
+        assert!(pack(&[cue(0.0, 1.0, "  ")], &Budgets::default()).is_empty());
+    }
+}

crates/cuepack/src/main.rs

@@ -0,0 +1,111 @@
+//! Packs cues into passages worth embedding.
+
+use std::io::{BufWriter, Read, Write};
+use std::path::PathBuf;
+use std::process::ExitCode;
+
+use clap::Parser;
+use corpus_core::{exit, ndjson};
+use cuekit::Cue;
+use cuepack::Budgets;
+
+/// Pack a cue stream into passages, one chunk record per line.
+///
+/// Takes either the cue stream `cuenorm` writes or a subtitle file directly.
+/// Each passage keeps the start of its first cue and the end of its last, so
+/// a result can be cited to the second — the same record a book's chunks use,
+/// with times where they have pages.
+#[derive(Parser)]
+#[command(
+    version,
+    about,
+    after_help = "\
+Example:
+
+    cuenorm episode.vtt | cuepack --seconds 75 --max-tokens 512 > passages.ndjson"
+)]
+struct Args {
+    /// Cue stream or subtitle file to read. Reads stdin when absent.
+    path: Option<PathBuf>,
+
+    /// Flush a passage once it covers this much audio.
+    #[arg(long, default_value_t = Budgets::default().seconds, value_name = "SECONDS")]
+    seconds: f64,
+
+    /// Never let a passage exceed this many tokens.
+    #[arg(long, default_value_t = Budgets::default().max_tokens, value_name = "TOKENS")]
+    max_tokens: usize,
+}
+
+fn main() -> ExitCode {
+    let args = Args::parse();
+    let budgets = Budgets {
+        seconds: args.seconds,
+        max_tokens: args.max_tokens,
+    };
+    if !budgets.seconds.is_finite() || budgets.seconds <= 0.0 || budgets.max_tokens == 0 {
+        eprintln!("cuepack: --seconds and --max-tokens must be above zero");
+        return ExitCode::from(exit::USAGE);
+    }
+
+    let content = match read(args.path.as_deref()) {
+        Ok(content) => content,
+        Err(e) => {
+            eprintln!("cuepack: {e}");
+            return ExitCode::from(exit::IO);
+        }
+    };
+
+    let cues = match cues(&content) {
+        Ok(cues) => cues,
+        Err(e) => {
+            eprintln!("cuepack: {e}");
+            return ExitCode::from(exit::DATA);
+        }
+    };
+
+    let mut out = BufWriter::new(std::io::stdout().lock());
+    for passage in cuepack::pack(&cues, &budgets) {
+        if let Err(e) = ndjson::write_line(&mut out, &passage) {
+            // A consumer that stopped reading is not a failure.
+            if e.is_broken_pipe() {
+                return ExitCode::from(exit::OK);
+            }
+            eprintln!("cuepack: {e}");
+            return ExitCode::from(exit::IO);
+        }
+    }
+    match out.flush() {
+        Ok(()) => ExitCode::from(exit::OK),
+        Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => ExitCode::from(exit::OK),
+        Err(e) => {
+            eprintln!("cuepack: {e}");
+            ExitCode::from(exit::IO)
+        }
+    }
+}
+
+/// Read cues from either encoding: a cue stream begins with a JSON object
+/// and a subtitle file does not.
+fn cues(content: &str) -> Result<Vec<Cue>, ndjson::Error> {
+    let first = content
+        .lines()
+        .find(|line| !line.trim().is_empty())
+        .unwrap_or("");
+    if first.trim_start().starts_with('{') {
+        return ndjson::read(content.as_bytes()).collect();
+    }
+    Ok(cuekit::parse(content))
+}
+
+fn read(path: Option<&std::path::Path>) -> std::io::Result<String> {
+    let bytes = match path {
+        Some(path) => std::fs::read(path)?,
+        None => {
+            let mut bytes = Vec::new();
+            std::io::stdin().read_to_end(&mut bytes)?;
+            bytes
+        }
+    };
+    Ok(String::from_utf8_lossy(&bytes).into_owned())
+}