more segementing simplication

Author: yujonglee

apps/desktop/src/utils/segment.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "vitest";
-import { buildSegments, type SegmentWord } from "./segment";
+import { buildSegments } from "./segment";
 
 describe("buildSegments", () => {
   const testCases = [
@@ -82,10 +82,10 @@ describe("buildSegments", () => {
       ],
     },
     {
-      name: "does not merge speaker turns once the max gap is reached",
+      name: "does not merge speaker turns once it exceeds the max gap",
       finalWords: [
         { text: "first", start_ms: 0, end_ms: 100, channel: 0 },
-        { text: "return", start_ms: 2100, end_ms: 2200, channel: 0 },
+        { text: "return", start_ms: 2101, end_ms: 2201, channel: 0 },
         { text: "other", start_ms: 150, end_ms: 200, channel: 1 },
       ],
       partialWords: [],
@@ -104,25 +104,6 @@ describe("buildSegments", () => {
         }),
       ],
     },
-    {
-      name: "applies the provided transform to every word",
-      finalWords: [
-        { text: "final", start_ms: 0, end_ms: 100, channel: 0 },
-      ],
-      partialWords: [
-        { text: "partial", start_ms: 120, end_ms: 160, channel: 0 },
-      ],
-      transform: (word: SegmentWord) => ({ ...word, label: word.isFinal ? "final" : "partial" }),
-      expected: [
-        expect.objectContaining({
-          key: { channel: 0 },
-          words: [
-            expect.objectContaining({ text: "final", label: "final" }),
-            expect.objectContaining({ text: "partial", label: "partial" }),
-          ],
-        }),
-      ],
-    },
     {
       name: "merges when gap is exactly at maxGapMs threshold (2000ms)",
       finalWords: [
@@ -197,8 +178,8 @@ describe("buildSegments", () => {
     },
   ];
 
-  test.each(testCases)("$name", ({ finalWords, partialWords, transform, expected }) => {
-    const segments = buildSegments(finalWords, partialWords, transform);
+  test.each(testCases)("$name", ({ finalWords, partialWords, expected }) => {
+    const segments = buildSegments(finalWords, partialWords);
     expect(segments).toEqual(expected);
   });
 });

apps/desktop/src/utils/segment.ts

@@ -19,175 +19,63 @@ type SegmentKey = { channel: number } | { speaker_index: number; channel: number
 export function buildSegments<
   TFinal extends WordLike,
   TPartial extends WordLike,
-  TWord extends SegmentWord = SegmentWord,
 >(
   finalWords: readonly TFinal[],
   partialWords: readonly TPartial[],
-  transform?: (word: SegmentWord) => TWord,
-): Segment<TWord>[] {
-  const mapWord = transform ?? ((word) => word as TWord);
-  const wordsByChannel = groupWordsByChannel(finalWords, partialWords, mapWord);
-  return createSpeakerTurns(wordsByChannel);
-}
-
-function toSegmentWord(word: WordLike, isFinal: boolean): SegmentWord {
-  return {
-    text: word.text,
-    start_ms: word.start_ms,
-    end_ms: word.end_ms,
-    channel: word.channel,
-    isFinal,
-  };
-}
-
-function addWordToChannel<TWord extends SegmentWord>(
-  channels: Map<number, TWord[]>,
-  word: TWord,
-): void {
-  const channelWords = channels.get(word.channel) ?? [];
-  channelWords.push(word);
-  channels.set(word.channel, channelWords);
-}
-
-function groupWordsByChannel<
-  TFinal extends WordLike,
-  TPartial extends WordLike,
-  TWord extends SegmentWord,
->(
-  finalWords: readonly TFinal[],
-  partialWords: readonly TPartial[],
-  mapWord: (word: SegmentWord) => TWord,
-): Map<number, TWord[]> {
-  const channels = new Map<number, TWord[]>();
-
-  for (const word of finalWords) {
-    addWordToChannel(channels, mapWord(toSegmentWord(word, true)));
-  }
-
-  for (const word of partialWords) {
-    addWordToChannel(channels, mapWord(toSegmentWord(word, false)));
-  }
-
-  for (const words of channels.values()) {
-    words.sort((a, b) => a.start_ms - b.start_ms);
-  }
-
-  return channels;
+): Segment[] {
+  const allWords: SegmentWord[] = [
+    ...finalWords.map((word) => ({
+      text: word.text,
+      start_ms: word.start_ms,
+      end_ms: word.end_ms,
+      channel: word.channel,
+      isFinal: true,
+    })),
+    ...partialWords.map((word) => ({
+      text: word.text,
+      start_ms: word.start_ms,
+      end_ms: word.end_ms,
+      channel: word.channel,
+      isFinal: false,
+    })),
+  ];
+
+  return createSpeakerTurns(allWords);
 }
 
-function flattenAndSortWords<TWord extends SegmentWord>(
-  wordsByChannel: Map<number, TWord[]>,
-): TWord[] {
-  const allWords: TWord[] = [];
-  wordsByChannel.forEach((words) => allWords.push(...words));
-  allWords.sort((a, b) => a.start_ms - b.start_ms);
-  return allWords;
-}
-
-function splitIntoInitialTurns<TWord extends SegmentWord>(
-  sortedWords: TWord[],
+function createSpeakerTurns<TWord extends SegmentWord>(
+  words: TWord[],
+  maxGapMs = 2000,
 ): Segment<TWord>[] {
-  if (sortedWords.length === 0) {
+  if (words.length === 0) {
     return [];
   }
 
-  const turns: Segment<TWord>[] = [];
-  let currentTurn: Segment<TWord> = {
-    key: { channel: sortedWords[0].channel },
-    words: [sortedWords[0]],
-  };
+  const sortedWords = [...words].sort((a, b) => a.start_ms - b.start_ms);
+  const segments: Segment<TWord>[] = [];
+  const currentByChannel = new Map<number, Segment<TWord>>();
 
-  for (let i = 1; i < sortedWords.length; i++) {
-    const word = sortedWords[i];
+  for (const word of sortedWords) {
+    const current = currentByChannel.get(word.channel);
 
-    if (word.channel === currentTurn.key.channel) {
-      currentTurn.words.push(word);
-    } else {
-      turns.push(currentTurn);
-      currentTurn = { key: { channel: word.channel }, words: [word] };
+    if (!current) {
+      const newSegment = { key: { channel: word.channel }, words: [word] };
+      currentByChannel.set(word.channel, newSegment);
+      segments.push(newSegment);
+      continue;
     }
-  }
-
-  turns.push(currentTurn);
-  return turns;
-}
-
-function groupSegmentsByChannel<TWord extends SegmentWord>(
-  segments: Segment<TWord>[],
-): Map<number, Segment<TWord>[]> {
-  const byChannel = new Map<number, Segment<TWord>[]>();
-
-  for (const segment of segments) {
-    const channelSegments = byChannel.get(segment.key.channel) ?? [];
-    channelSegments.push(segment);
-    byChannel.set(segment.key.channel, channelSegments);
-  }
-
-  return byChannel;
-}
-
-function getSegmentStartTime<TWord extends SegmentWord>(
-  segment: Segment<TWord>,
-): number {
-  return segment.words[0]?.start_ms ?? 0;
-}
-
-function calculateTimingGap<TWord extends SegmentWord>(
-  firstSegment: Segment<TWord>,
-  secondSegment: Segment<TWord>,
-): number {
-  if (firstSegment.words.length === 0 || secondSegment.words.length === 0) {
-    return Infinity;
-  }
-
-  const lastWordOfFirst = firstSegment.words[firstSegment.words.length - 1];
-  const firstWordOfSecond = secondSegment.words[0];
-  return firstWordOfSecond.start_ms - lastWordOfFirst.end_ms;
-}
-
-function mergeSegmentsByGap<TWord extends SegmentWord>(
-  segments: Segment<TWord>[],
-  channel: number,
-  maxGapMs: number,
-): Segment<TWord>[] {
-  segments.sort((a, b) => getSegmentStartTime(a) - getSegmentStartTime(b));
-
-  const merged: Segment<TWord>[] = [];
-  let currentMerged = { key: { channel }, words: [...segments[0].words] };
 
-  for (let i = 1; i < segments.length; i++) {
-    const nextSegment = segments[i];
-    const gap = calculateTimingGap(currentMerged, nextSegment);
+    const lastWord = current.words[current.words.length - 1];
+    const gap = word.start_ms - lastWord.end_ms;
 
-    if (gap < maxGapMs) {
-      currentMerged.words.push(...nextSegment.words);
+    if (gap <= maxGapMs) {
+      current.words.push(word);
     } else {
-      merged.push(currentMerged);
-      currentMerged = { key: { channel }, words: [...nextSegment.words] };
+      const newSegment = { key: { channel: word.channel }, words: [word] };
+      currentByChannel.set(word.channel, newSegment);
+      segments.push(newSegment);
     }
   }
 
-  merged.push(currentMerged);
-  return merged;
-}
-
-function createSpeakerTurns<TWord extends SegmentWord>(
-  wordsByChannel: Map<number, TWord[]>,
-  maxGapMs = 2000,
-): Segment<TWord>[] {
-  const sortedWords = flattenAndSortWords(wordsByChannel);
-  if (sortedWords.length === 0) {
-    return [];
-  }
-
-  const initialTurns = splitIntoInitialTurns(sortedWords);
-  const turnsByChannel = groupSegmentsByChannel(initialTurns);
-
-  const finalSegments: Segment<TWord>[] = [];
-  for (const [channel, channelTurns] of turnsByChannel) {
-    finalSegments.push(...mergeSegmentsByGap(channelTurns, channel, maxGapMs));
-  }
-
-  finalSegments.sort((a, b) => getSegmentStartTime(a) - getSegmentStartTime(b));
-  return finalSegments;
+  return segments;
 }