cueseg: chapter by vectors somebody else made, where there are some

The hashed bag of words asks only whether two stretches repeat the same
words, and cannot tell that ships and vessels are one subject. An embedding
can. segment_with takes one vector per window and falls back to the words
when it is given none, or the wrong number of them — a caller that embedded
some other recording's windows is asking about that recording.

windows_of gives the text to embed, in the order the vectors come back in,
so nobody has to reimplement the windowing pass to use this.

Bound in all three languages, and on the command line as --windows and
--embeddings, which pipe through emb. The C ABI version goes to 2:
catena_segment_cues gained an argument.

crates/bindings/src/capi.rs

@@ -30,7 +30,9 @@ use crate::{ChunkBudgets, CueBudgets};
 
 /// Bumped when the meaning or signature of anything here changes. Additions
 /// leave it alone.
-const ABI_VERSION: c_uint = 1;
+///
+/// 2: `catena_segment_cues` gained an embeddings argument.
+const ABI_VERSION: c_uint = 2;
 
 thread_local! {
     static LAST_ERROR: RefCell<Option<CString>> = const { RefCell::new(None) };
@@ -271,6 +273,7 @@ pub unsafe extern "C" fn catena_pack_cues(
 pub unsafe extern "C" fn catena_segment_cues(
     cues: *const c_char,
     duration: c_double,
+    embeddings: *const c_char,
 ) -> *mut c_char {
     clear_error();
     let Some(cues) = (unsafe { borrow(cues, "cues") }) else {
@@ -279,7 +282,42 @@ pub unsafe extern "C" fn catena_segment_cues(
     let Some(cues) = take_json::<Vec<cuekit::Cue>>(cues, "cues") else {
         return ptr::null_mut();
     };
-    give_json(&crate::segment_cues(&cues, duration))
+
+    let vectors: Option<Vec<Vec<f64>>> = match embeddings.is_null() {
+        true => None,
+        false => {
+            let Some(text) = (unsafe { borrow(embeddings, "embeddings") }) else {
+                return ptr::null_mut();
+            };
+            match take_json::<Vec<Vec<f64>>>(text, "embeddings") {
+                Some(vectors) => Some(vectors),
+                None => return ptr::null_mut(),
+            }
+        }
+    };
+    give_json(&crate::segment_cues_with(
+        &cues,
+        duration,
+        vectors.as_deref(),
+    ))
+}
+
+/// The text of each window a transcript is measured in, to embed. Returns a
+/// JSON array of strings.
+///
+/// # Safety
+///
+/// `cues` must be NUL-terminated UTF-8.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn catena_cue_windows(cues: *const c_char) -> *mut c_char {
+    clear_error();
+    let Some(cues) = (unsafe { borrow(cues, "cues") }) else {
+        return ptr::null_mut();
+    };
+    let Some(cues) = take_json::<Vec<cuekit::Cue>>(cues, "cues") else {
+        return ptr::null_mut();
+    };
+    give_json(&crate::cue_windows(&cues))
 }
 
 /// Split a collection into the works inside it. Returns a JSON array of
@@ -550,6 +588,6 @@ mod tests {
 
     #[test]
     fn the_abi_version_is_the_one_the_header_declares() {
-        assert_eq!(catena_abi_version(), 1);
+        assert_eq!(catena_abi_version(), 2);
     }
 }

crates/bindings/src/lib.rs

@@ -133,11 +133,32 @@ pub fn pack_cues(cues: &[Cue], budgets: CueBudgets) -> Vec<Chunk> {
 /// `duration` is how long the recording is, in seconds; it decides how many
 /// chapters to aim for. Zero or less means take it from the last cue.
 pub fn segment_cues(cues: &[Cue], duration: f64) -> Vec<cueseg::Segment> {
+    segment_cues_with(cues, duration, None)
+}
+
+/// The same, given vectors somebody else computed for the windows.
+///
+/// Without them the windows are compared by the words they repeat, which
+/// cannot tell that "ships" and "vessels" are one subject. [`cue_windows`]
+/// gives the text to embed, in the order the vectors must come back in.
+pub fn segment_cues_with(
+    cues: &[Cue],
+    duration: f64,
+    embeddings: Option<&[Vec<f64>]>,
+) -> Vec<cueseg::Segment> {
     let duration = match duration > 0.0 {
         true => duration,
         false => cues.last().map_or(0.0, |cue| cue.end),
     };
-    cueseg::segment(cues, duration)
+    cueseg::segment_with(cues, duration, embeddings)
+}
+
+/// The text of each window a transcript is measured in, to embed.
+pub fn cue_windows(cues: &[Cue]) -> Vec<String> {
+    cueseg::windows_of(cues)
+        .into_iter()
+        .map(|window| window.text)
+        .collect()
 }
 
 /// Split a collection into the works inside it.
@@ -236,24 +257,37 @@ mod tests {
 
     #[test]
     fn a_transcript_comes_back_as_chapters() {
-        // Long enough to be worth chaptering at all: the segmenter refuses
-        // to cut a recording into pieces shorter than its minimum.
+        // Two halves about two different things, so the words themselves say
+        // where the boundary is.
         let cues: Vec<Cue> = (0..1200)
             .map(|i| Cue {
                 v: corpus_core::V,
                 start: i as f64 * 5.0,
                 end: i as f64 * 5.0 + 5.0,
-                text: match i % 2 {
-                    0 => "the mill stood at the northern end of the valley".into(),
-                    _ => "parliament dissolved itself again amid objections".into(),
+                text: match i < 600 {
+                    true => "the mill stood at the northern end of the valley".into(),
+                    false => "parliament dissolved itself again amid objections".into(),
                 },
             })
             .collect();
+
         let chapters = segment_cues(&cues, 6000.0);
         assert!(!chapters.is_empty());
         assert_eq!(chapters, cueseg::segment(&cues, 6000.0));
         // The duration can be left to the cues themselves.
-        assert_eq!(segment_cues(&cues, 0.0), cueseg::segment(&cues, 6000.0));
+        assert_eq!(segment_cues(&cues, 0.0).len(), chapters.len());
+
+        // The windows to embed are the ones the segmenter measures, and
+        // vectors given for them are used rather than the words.
+        let windows = cue_windows(&cues);
+        assert_eq!(windows.len(), cueseg::windows_of(&cues).len());
+        let flat: Vec<Vec<f64>> = windows.iter().map(|_| vec![1.0, 0.0]).collect();
+        let with_flat = segment_cues_with(&cues, 6000.0, Some(&flat));
+        assert_ne!(
+            with_flat.len(),
+            chapters.len(),
+            "vectors saying every window is alike gave the same chapters as the words"
+        );
     }
 
     #[test]

crates/bindings/src/python.rs

@@ -162,15 +162,28 @@ fn pack_cues<'py>(
 ///
 /// `duration` is the recording's length in seconds; omit it and the last cue
 /// decides.
+/// `embeddings` is one vector per window, in the order `cue_windows` gives
+/// them; without it the windows are compared by the words they repeat.
 #[pyfunction]
-#[pyo3(signature = (cues, *, duration = None))]
+#[pyo3(signature = (cues, *, duration = None, embeddings = None))]
 fn segment_cues<'py>(
     py: Python<'py>,
     cues: &Bound<'py, PyAny>,
     duration: Option<f64>,
+    embeddings: Option<Vec<Vec<f64>>>,
 ) -> PyResult<Bound<'py, PyList>> {
     let cues: Vec<cuekit::Cue> = from_py_list(cues)?;
-    list_to_py(py, &crate::segment_cues(&cues, duration.unwrap_or(0.0)))
+    list_to_py(
+        py,
+        &crate::segment_cues_with(&cues, duration.unwrap_or(0.0), embeddings.as_deref()),
+    )
+}
+
+/// The text of each window a transcript is measured in, to embed.
+#[pyfunction]
+fn cue_windows(cues: &Bound<'_, PyAny>) -> PyResult<Vec<String>> {
+    let cues: Vec<cuekit::Cue> = from_py_list(cues)?;
+    Ok(crate::cue_windows(&cues))
 }
 
 /// Split a collection into the works inside it.
@@ -244,6 +257,7 @@ fn _catena(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_function(wrap_pyfunction!(parse_subtitles, module)?)?;
     module.add_function(wrap_pyfunction!(pack_cues, module)?)?;
     module.add_function(wrap_pyfunction!(segment_cues, module)?)?;
+    module.add_function(wrap_pyfunction!(cue_windows, module)?)?;
     module.add_function(wrap_pyfunction!(split_collection, module)?)?;
     module.add_function(wrap_pyfunction!(fuse_hits, module)?)?;
     module.add_function(wrap_pyfunction!(identify, module)?)?;

crates/bindings/src/wasm.rs

@@ -212,10 +212,31 @@ pub fn pack_cues(
 ///
 /// `duration` is the recording's length in seconds; omit it and the last cue
 /// decides.
+/// `embeddings` is one vector per window, in the order `cueWindows` gives
+/// them; without it the windows are compared by the words they repeat.
 #[wasm_bindgen(js_name = segmentCues, skip_typescript)]
-pub fn segment_cues(cues: JsValue, duration: Option<f64>) -> Result<JsValue, JsValue> {
+pub fn segment_cues(
+    cues: JsValue,
+    duration: Option<f64>,
+    embeddings: JsValue,
+) -> Result<JsValue, JsValue> {
     let cues: Vec<cuekit::Cue> = from_js(cues, "cues")?;
-    to_js(&crate::segment_cues(&cues, duration.unwrap_or(0.0)))
+    let vectors: Option<Vec<Vec<f64>>> = match embeddings.is_undefined() || embeddings.is_null() {
+        true => None,
+        false => Some(from_js(embeddings, "embeddings")?),
+    };
+    to_js(&crate::segment_cues_with(
+        &cues,
+        duration.unwrap_or(0.0),
+        vectors.as_deref(),
+    ))
+}
+
+/// The text of each window a transcript is measured in, to embed.
+#[wasm_bindgen(js_name = cueWindows, skip_typescript)]
+pub fn cue_windows(cues: JsValue) -> Result<JsValue, JsValue> {
+    let cues: Vec<cuekit::Cue> = from_js(cues, "cues")?;
+    to_js(&crate::cue_windows(&cues))
 }
 
 /// Split a collection into the works inside it.
@@ -288,7 +309,8 @@ const SIGNATURES: &'static str = r#"
 export function chunkMarkdown(markdown: string, target?: number, overlap?: number, min_tokens?: number, hard_max?: number): Chunk[];
 export function parseSubtitles(content: string, keep_repeats?: boolean): Cue[];
 export function packCues(cues: Cue[], seconds?: number, max_tokens?: number): Chunk[];
-export function segmentCues(cues: Cue[], duration?: number): Segment[];
+export function segmentCues(cues: Cue[], duration?: number, embeddings?: number[][]): Segment[];
+export function cueWindows(cues: Cue[]): string[];
 export function splitCollection(markdown: string): Section[];
 export function fuseHits(hits: Hit[], k?: number, weights?: Record<string, number>, limit?: number): Hit[];
 export function identify(name: string, data: Uint8Array): Doc;