Repeat guard cuts at silence gaps; trim opening silence

README.md

@@ -47,8 +47,9 @@ LFM3-Omni-2B often re-asks its last question ("Red is a great choice. What's you
 favorite color?"), including with the original Space's exact prompt; rules, examples
 and per-turn notes in the prompt didn't stop it. `public/workers/repeat-guard.js`
 catches a repeated question or sentence in the text stream, which runs ahead of
-playback, then ends the reply just before it and keeps it out of the history. It can't
-cut a reply's first sentence, or audio already heard, and it can't make her lead better.
+playback, then ends the reply at the silence after the last good sentence and keeps
+the repeat out of subtitles and history. It leaves a reply's first sentence alone
+(silence would be worse), and it can't make her lead better.
 
 ## Run
 

index.html

@@ -12,7 +12,7 @@
       rel="stylesheet"
     />
     
-  <link rel="stylesheet" crossorigin href="./chunk-sarxdv2b.css"><script type="module" crossorigin src="./chunk-qht2aa6q.js"></script></head>
+  <link rel="stylesheet" crossorigin href="./chunk-sarxdv2b.css"><script type="module" crossorigin src="./chunk-zg54109f.js"></script></head>
   <body>
     <main id="app">
       <!-- The avatar is the hero: a full-viewport 3D stage. -->

workers/model-worker.js

@@ -15,7 +15,7 @@
  *   { kind: "progress", loaded?, total?, fraction?, status? }
  *   { kind: "ready", device, sampleRate }   { kind: "failed", name, message }
  *   { kind: "text", id, text }     { kind: "audio", id, samples }
- *   { kind: "truncate", id, samples, kept }   the reply ends at this audio sample
+ *   { kind: "truncate", id, samples }   the reply ends at this audio sample
  *   { kind: "finished", id, reason, usage, generated, error? }
  */
 
@@ -24,7 +24,7 @@ import {
   ContextOverflowError,
   DeviceLostError,
 } from "../vendor/Lfm3OmniForConditionalGeneration.js";
-import { RepeatGuard } from "./repeat-guard.js";
+import { RepeatGuard, findCut, firstSound } from "./repeat-guard.js";
 
 const LOAD_ATTEMPTS = 3; // Chromium occasionally returns no adapter on a cold GPU process
 const LOAD_RETRY_MS = 1500;
@@ -33,6 +33,7 @@ const LOAD_RETRY_MS = 1500;
 const SEARCH = new URL(location.href).searchParams;
 const HISTORY_BUDGET_TOKENS = Number(SEARCH.get("historyBudget")) || 3000;
 
+const PRE_SPEECH_SAMPLES = 1200; // keep 50 ms before her first sound
 /** @type {any} */ let model = null;
 /** @type {any} */ let session = null;
 let system = "";
@@ -118,50 +119,72 @@ async function respond({ id, audio }) {
     history.push({ role: "user", content: [{ type: "audio", audio }] });
 
     guard.beginReply();
-    /** @type {string | null} */ let keptOnRepeat = null;
+    let repeatFound = false;
+    let truncated = false;
+    /** @type {Float32Array[]} */ const heard = [];
     for (;;) {
       const stream = session.generate(messages(), { generation_mode: "interleaved", signal });
       let emitted = false;
       const started = performance.now();
       try {
         for await (const ev of stream) {
-          if (ev.type === "text") {
+          if (ev.type === "text" && !repeatFound) {
             emitted = true;
-            post({ kind: "text", id, text: ev.text });
-            const repeat = guard.push(ev.text, generated);
-            if (repeat) {
-              keptOnRepeat = repeat.kept;
-              post({ kind: "truncate", id, samples: repeat.cutAt, kept: repeat.kept });
-              controller.abort();
-              break;
-            }
+            const { released, repeat } = guard.push(ev.text);
+            if (released) post({ kind: "text", id, text: released });
+            repeatFound = repeat;
           } else if (ev.type === "audio") {
             emitted = true;
-            generated += ev.audio.length;
-            const samples = ev.audio.slice();
+            let samples = ev.audio.slice();
+            if (generated === 0) {
+              // Her replies open with ~0.8 s of silence. Don't make anyone wait for it.
+              const start = firstSound(samples);
+              if (start < 0) continue;
+              samples = samples.slice(Math.max(0, start - PRE_SPEECH_SAMPLES));
+            }
+            generated += samples.length;
+            if (repeatFound || guard.kept) heard.push(samples.slice());
             if (simulateRtf) {
               const due = started + (generated / model.audio_sampling_rate / simulateRtf) * 1000;
               await new Promise((r) => setTimeout(r, Math.max(0, due - performance.now())));
             }
             post({ kind: "audio", id, samples }, [samples.buffer]);
+            if (repeatFound) {
+              const audio = new Float32Array(generated);
+              let at = 0;
+              for (const h of heard) {
+                audio.set(h, at);
+                at += h.length;
+              }
+              const cut = findCut(audio, model.audio_sampling_rate, guard.kept.length);
+              if (cut !== null) {
+                post({ kind: "truncate", id, samples: cut });
+                truncated = true;
+                controller.abort();
+                break;
+              }
+            }
           }
         }
       } catch (err) {
-        if (!emitted && err instanceof ContextOverflowError && dropOldestExchange()) continue;
-        throw err;
+        if (!truncated) {
+          if (!emitted && err instanceof ContextOverflowError && dropOldestExchange()) continue;
+          throw err;
+        }
       }
-      const result = await stream.result;
-      usage = result.usage;
-      if (keptOnRepeat !== null) {
+      const result = await stream.result.catch(() => null);
+      usage = result?.usage ?? usage;
+      reason = result?.finish_reason ?? "stop";
+      if (repeatFound) {
         // History holds what she said, not the repeat, or it would teach her to loop.
+        history.push({ role: "assistant", content: guard.kept });
         reason = "stop";
-        history.push({ role: "assistant", content: keptOnRepeat });
-        guard.endReply();
       } else {
-        reason = result.finish_reason;
+        const rest = guard.flush();
+        if (rest) post({ kind: "text", id, text: rest });
         history.push(result.message);
-        if (reason === "stop") guard.endReply();
       }
+      guard.endReply();
       break;
     }
   } catch (err) {

workers/repeat-guard.js

@@ -3,8 +3,10 @@
  * LFM3-Omni-2B tends to re-ask its own last question ("Red is a great choice.
  * What's your favorite color?"), whatever the prompt says. Speech generates
  * about 3x faster than it plays, so a repeat is usually visible in the text
- * before it is heard. This spots it, and says where in the audio its sentence
- * began so the reply can end cleanly just before it. Pure logic, for testing.
+ * before it is heard. This spots it in the text, and finds the silence that ends
+ * the last good sentence in the audio, so the reply can stop cleanly there. Text
+ * runs several words ahead of its audio, so the two can't be matched sample for
+ * sample. Pure logic, for testing.
  */
 
 const STOP_WORDS = new Set(
@@ -41,32 +43,49 @@ export class RepeatGuard {
 
   beginReply() {
     this._current = "";
-    this._currentStart = 0;
     this._kept = "";
     /** @type {string[]} */
     this._sentences = [];
   }
 
   /**
-   * Feed one text delta, with how many audio samples existed when it arrived.
-   * @param {string} text @param {number} samples
-   * @returns {{ cutAt: number, kept: string } | null} set when a repeat was found
+   * Feed one text delta. Text is released a sentence at a time, once it is known
+   * not to be a repeat, so a repeat never reaches the subtitles.
+   * @param {string} text
+   * @returns {{ released: string, repeat: boolean }}
    */
-  push(text, samples) {
+  push(text) {
+    let released = "";
     for (let i = 0; i < text.length; i++) {
       const ch = text[i];
-      if (!this._current.trim()) this._currentStart = samples;
       this._current += ch;
       const next = text[i + 1];
       const closes = ch === "?" || ch === "!" || (ch === "." && (next === undefined || /\s/.test(next)));
       if (!closes) continue;
       const sentence = this._current.trim();
+      if (this._isRepeat(sentence)) return { released, repeat: true };
+      released += this._current;
       this._current = "";
-      if (this._isRepeat(sentence)) return { cutAt: this._currentStart, kept: this._kept.trim() };
       this._kept += ` ${sentence}`;
       this._sentences.push(sentence);
     }
-    return null;
+    return { released, repeat: false };
+  }
+
+  /** What's left unreleased when the reply ends normally. */
+  flush() {
+    const rest = this._current;
+    const sentence = rest.trim();
+    this._current = "";
+    if (sentence) {
+      this._kept += ` ${sentence}`;
+      this._sentences.push(sentence);
+    }
+    return rest;
+  }
+
+  get kept() {
+    return this._kept.trim();
   }
 
   /** The reply is over: remember what was actually said. */
@@ -84,3 +103,43 @@ export class RepeatGuard {
     return this.said.some((s) => s.endsWith("?") === isQuestion && similar(s, sentence));
   }
 }
+
+const SPEECH_CHARS_PER_SEC = 15;
+const GAP_RMS = 0.004;
+const GAP_FRAME_SEC = 0.02;
+const MIN_GAP_SEC = 0.3;
+
+/**
+ * Where to end a reply whose good part is `keptChars` of text. Looks for the
+ * first long silence around where that much speech should have finished; commas
+ * make shorter pauses earlier on. Null until enough audio exists to decide.
+ * @param {Float32Array} audio speech so far, starting at its first sound
+ * @param {number} sampleRate @param {number} keptChars
+ * @returns {number | null} sample to stop at
+ */
+export function findCut(audio, sampleRate, keptChars) {
+  const estimate = keptChars / SPEECH_CHARS_PER_SEC;
+  const hop = Math.round(GAP_FRAME_SEC * sampleRate);
+  const earliest = Math.round(estimate * 0.75 * sampleRate);
+  const giveUpAt = Math.round((estimate * 1.6 + 0.5) * sampleRate);
+  let quietFrom = -1;
+  for (let i = earliest; i + hop <= audio.length; i += hop) {
+    let sum = 0;
+    for (let j = i; j < i + hop; j++) sum += audio[j] * audio[j];
+    const quiet = Math.sqrt(sum / hop) < GAP_RMS;
+    if (!quiet) quietFrom = -1;
+    else if (quietFrom < 0) quietFrom = i;
+    if (quietFrom >= 0 && i + hop - quietFrom >= MIN_GAP_SEC * sampleRate) return quietFrom + Math.round(0.1 * sampleRate);
+    if (i >= giveUpAt) return giveUpAt;
+  }
+  return null;
+}
+
+/**
+ * Index of a reply's first sound, so she doesn't open with silence.
+ * @param {Float32Array} audio @returns {number} -1 if it's all silence so far
+ */
+export function firstSound(audio) {
+  for (let i = 0; i < audio.length; i++) if (Math.abs(audio[i]) > 0.01) return i;
+  return -1;
+}