diff --git a/CLAUDE.md b/CLAUDE.md index 32712ed..b89f6cb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -70,6 +70,7 @@ podwhisperer/ │ ├── index.ts # Lambda handler (durable workflow) │ ├── types.ts # Shared types │ ├── steps/ # Pipeline step implementations +│ │ ├── fill-timing-gaps.ts │ │ ├── llm-refinement.ts │ │ ├── replacement.ts │ │ └── segments-normalization.ts @@ -100,7 +101,7 @@ podwhisperer/ 2. **Pipeline Lambda** (durable execution) orchestrates the workflow: - Sends job to SQS queue - Waits for callback from GPU worker - - Applies post-processing steps (replacement, LLM, normalization) + - Applies post-processing steps (fill timing gaps, replacement, LLM, normalization) - Generates captions (VTT, SRT, JSON) - Sends EventBridge notification 3. **GPU Worker** (ECS container on Managed Instances): @@ -124,7 +125,7 @@ podwhisperer/ **Configuration** (`packages/config/index.ts`): - Zod schemas for all pipeline configuration -- Exported types: `PipelineConfig`, `TranscriptionConfig`, etc. +- Exported types: `PipelineConfig`, `TranscriptionConfig`, `FillTimingGapsConfig`, etc. - Used by both CDK (at synth time) and Lambda (at runtime) ## Important Implementation Notes diff --git a/README.md b/README.md index 2940777..ea7f408 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,7 @@ flowchart TD CheckExists{Output Exists?} SQS[Send to SQS] WaitCallback[Wait for Callback] + FillGaps[Fill Timing Gaps] Replacement[Replacement Rules] LLM[LLM Refinement] Normalize[Segments Normalization] @@ -67,10 +68,10 @@ flowchart TD S3Upload --> EB --> CheckExists CheckExists -->|No| SQS --> Poll - CheckExists -->|Yes| Replacement + CheckExists -->|Yes| FillGaps Poll --> Download --> Transcribe --> Align --> Diarize --> Upload --> Callback Upload --> Raw - Callback --> WaitCallback --> Replacement --> LLM --> Normalize --> Finalize --> Captions --> Notify + Callback --> WaitCallback --> FillGaps --> Replacement --> LLM --> Normalize --> Finalize --> Captions --> Notify Finalize --> Refined Captions --> VTT & SRT & JSON ``` @@ -88,7 +89,15 @@ The WhisperX container runs on GPU-enabled ECS Managed Instances: 5. **Diarization** - Identifies and labels different speakers 6. **Upload** - Saves raw transcript to S3 `output/` prefix -### Step 1: Replacement Rules (Optional) +### Step 1: Fill Timing Gaps + +WhisperX sometimes produces words with missing `start`/`end` timestamps, particularly for words containing numbers or currency symbols (e.g., `"15"`, `"$8"`, `"$0.25"`). This step fills those gaps using character-proportional interpolation between neighboring anchor words, giving all downstream steps clean timing data. + +- Gaps are distributed proportionally based on word character length +- Padding is applied dynamically based on the segment's speech rate +- Filled words are marked with `score: 0` to distinguish them from original timestamps + +### Step 2: Replacement Rules (Optional) Apply regex or literal string replacements to fix common transcription errors: @@ -101,7 +110,7 @@ Apply regex or literal string replacements to fix common transcription errors: } ``` -### Step 2: LLM Refinement (Optional) +### Step 3: LLM Refinement (Optional) Uses Amazon Bedrock to improve transcript quality: @@ -109,7 +118,7 @@ Uses Amazon Bedrock to improve transcript quality: - Identifies speakers by name when possible (e.g., "SPEAKER_00" becomes "Luciano") - Validates suggestions to prevent aggressive rewrites -### Step 3: Segments Normalization +### Step 4: Segments Normalization Splits long transcript segments into caption-friendly chunks: @@ -117,7 +126,7 @@ Splits long transcript segments into caption-friendly chunks: - Splits at natural punctuation boundaries - Handles speaker changes within segments -### Step 4: Caption Generation +### Step 5: Caption Generation Generates multiple caption formats from the refined transcript: @@ -129,7 +138,7 @@ Optional features: - Word-by-word highlighting (underline, bold, or italic) - Speaker name prefixes (always, when speaker changes, or never) -### Step 5: Notification +### Step 6: Notification Sends an EventBridge event when the pipeline completes, including: @@ -270,6 +279,12 @@ The pipeline will automatically trigger and generate outputs in `s3://$BUCKET/ou | `skipIfOutputExists` | boolean | `false` | Skip transcription if output already exists | | `hfTokenSsmPath` | string | `"/podwhisperer/hf_token"` | SSM parameter path for HuggingFace token | +### Fill Timing Gaps + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `enabled` | boolean | `true` | Enable fill-timing-gaps step | + ### Replacement Rules Array of replacement rules, each with: diff --git a/lambdas/pipeline/index.ts b/lambdas/pipeline/index.ts index 7a701c8..2a39b80 100644 --- a/lambdas/pipeline/index.ts +++ b/lambdas/pipeline/index.ts @@ -20,6 +20,7 @@ import { import { SendMessageCommand, SQSClient } from '@aws-sdk/client-sqs' import { PipelineConfigSchema } from '@podwhisperer/config' import type { EventBridgeEvent } from 'aws-lambda' +import { fillTimingGaps } from './steps/fill-timing-gaps' import { llmRefinement } from './steps/llm-refinement' import { applyReplacements } from './steps/replacement' import { normalizeSegments } from './steps/segments-normalization' @@ -167,6 +168,42 @@ const handler = async (event: S3EventBridgeEvent, context: DurableContext) => { // refining steps + // Fill timing gaps — fills missing word timestamps before any other processing + await context.step('fill-timing-gaps', async () => { + if (!pipelineConfig.fillTimingGaps.enabled) { + context.logger.info('Fill timing gaps disabled, skipping step') + return + } + + // Load the transcript from S3 + const response = await s3.send( + new GetObjectCommand({ + Bucket: BUCKET_NAME, + Key: transcriptProcessingKey, + }), + ) + const bodyStr = await response.Body?.transformToString() + if (!bodyStr) { + throw new Error('Empty response body from S3') + } + const transcript = JSON.parse(bodyStr) as WhisperxResult + + // Fill timing gaps + const stats = fillTimingGaps(transcript, pipelineConfig.fillTimingGaps) + + // Save the updated transcript back to S3 + await s3.send( + new PutObjectCommand({ + Bucket: BUCKET_NAME, + Key: transcriptProcessingKey, + Body: JSON.stringify(transcript), + ContentType: 'application/json', + }), + ) + + context.logger.info('Filled word timing gaps', { stats }) + }) + // Replacement rules await context.step('replacement-rules', async () => { if (typeof pipelineConfig.replacementRules === 'undefined') { diff --git a/lambdas/pipeline/steps/fill-timing-gaps.test.ts b/lambdas/pipeline/steps/fill-timing-gaps.test.ts new file mode 100644 index 0000000..d318c6f --- /dev/null +++ b/lambdas/pipeline/steps/fill-timing-gaps.test.ts @@ -0,0 +1,497 @@ +import { describe, expect, it } from 'vitest' +import type { WhisperxResult, WhisperxWord } from '../types' +import { fillTimingGaps } from './fill-timing-gaps' + +/** Helper to create a word with timing */ +function w( + word: string, + start?: number, + end?: number, + score?: number | null, +): WhisperxWord { + const result: WhisperxWord = { word } + if (start !== undefined) result.start = start + if (end !== undefined) result.end = end + if (score !== undefined) result.score = score + return result +} + +/** Helper to create a simple transcript */ +function transcript( + segments: { + start: number + end: number + text: string + words?: WhisperxWord[] + }[], +): WhisperxResult { + return { segments } +} + +/** Helper to get words from a segment (avoids non-null assertions) */ +function getWords(t: WhisperxResult, segIdx = 0): WhisperxWord[] { + const words = t.segments[segIdx].words + expect(words).toBeDefined() + return words as WhisperxWord[] +} + +const enabledConfig = { enabled: true } + +describe('fillTimingGaps', () => { + it('fills a single missing word between two timed words', () => { + const t = transcript([ + { + start: 0, + end: 3, + text: 'hello world goodbye', + words: [w('hello', 0, 1, 0.9), w('world'), w('goodbye', 2, 3, 0.8)], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const filled = words[1] + expect(filled.start).toBeDefined() + expect(filled.end).toBeDefined() + expect(filled.score).toBeNull() + expect(filled.start).toBeGreaterThanOrEqual(1) + expect(filled.end).toBeLessThanOrEqual(2) + expect(stats.wordsFilled).toBe(1) + expect(stats.gapsFound).toBe(1) + expect(stats.gapTypes.middle).toBe(1) + }) + + it('fills multiple separate gaps in one segment', () => { + const t = transcript([ + { + start: 0, + end: 10, + text: 'aaa bbb ccc ddd eee', + words: [ + w('aaa', 0, 1, 0.9), + w('bbb'), + w('ccc', 4, 5, 0.8), + w('ddd'), + w('eee', 8, 10, 0.7), + ], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + expect(stats.gapsFound).toBe(2) + expect(stats.wordsFilled).toBe(2) + const words = getWords(t) + expect(words[1].start).toBeDefined() + expect(words[1].end).toBeDefined() + expect(words[3].start).toBeDefined() + expect(words[3].end).toBeDefined() + }) + + it('returns zero stats when no gaps exist', () => { + const t = transcript([ + { + start: 0, + end: 3, + text: 'hello world', + words: [w('hello', 0, 1, 0.9), w('world', 1, 3, 0.8)], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + expect(stats.wordsFilled).toBe(0) + expect(stats.gapsFound).toBe(0) + // Words should be unchanged + const words = getWords(t) + expect(words[0].score).toBe(0.9) + expect(words[1].score).toBe(0.8) + }) + + it('fills missing word at end of segment using segment.end as right anchor', () => { + const t = transcript([ + { + start: 0, + end: 5, + text: 'hello world', + words: [w('hello', 0, 2, 0.9), w('world')], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const filled = words[1] + expect(filled.start).toBeDefined() + expect(filled.end).toBeDefined() + expect(filled.end).toBeLessThanOrEqual(5) + expect(filled.score).toBeNull() + expect(stats.gapTypes.end).toBe(1) + }) + + it('fills missing word at start of segment using segment.start as left anchor', () => { + const t = transcript([ + { + start: 0, + end: 5, + text: 'hello world', + words: [w('hello'), w('world', 3, 5, 0.8)], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const filled = words[0] + expect(filled.start).toBeDefined() + expect(filled.end).toBeDefined() + expect(filled.start).toBeGreaterThanOrEqual(0) + expect(filled.score).toBeNull() + expect(stats.gapTypes.start).toBe(1) + }) + + it('distributes 3 consecutive missing words proportionally by character length', () => { + const t = transcript([ + { + start: 0, + end: 10, + text: 'start ab abcd abcdef end', + words: [ + w('start', 0, 1, 0.9), + w('ab'), // 2 chars + w('abcd'), // 4 chars + w('abcdef'), // 6 chars + w('end', 9, 10, 0.8), + ], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + expect(stats.wordsFilled).toBe(3) + expect(stats.gapsFound).toBe(1) + + const words = getWords(t) + // All 3 should have timing now + expect(words[1].start).toBeDefined() + expect(words[2].start).toBeDefined() + expect(words[3].start).toBeDefined() + + // Proportional: word with 6 chars should get ~3x the duration of word with 2 chars + const d1 = (words[1].end as number) - (words[1].start as number) + const d3 = (words[3].end as number) - (words[3].start as number) + expect(d3 / d1).toBeCloseTo(3, 0) + + // Continuity: each word starts where previous ends + expect(words[2].start).toBeCloseTo(words[1].end as number, 3) + expect(words[3].start).toBeCloseTo(words[2].end as number, 3) + }) + + it('fills all words when entire segment has missing timing', () => { + const t = transcript([ + { + start: 2, + end: 5, + text: 'hello world', + words: [w('hello'), w('world')], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + expect(stats.gapTypes.entireSegment).toBe(1) + expect(stats.wordsFilled).toBe(2) + + const words = getWords(t) + expect(words[0].start).toBeGreaterThanOrEqual(2) + expect(words[1].end).toBeLessThanOrEqual(5) + expect(words[0].score).toBeNull() + expect(words[1].score).toBeNull() + }) + + it('allocates time proportional to word length', () => { + // "hi" (2 chars) vs "conversation" (12 chars) — 1:6 ratio + const t = transcript([ + { + start: 0, + end: 14, + text: 'hi conversation', + words: [w('hi'), w('conversation')], + }, + ]) + + fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const d0 = (words[0].end as number) - (words[0].start as number) + const d1 = (words[1].end as number) - (words[1].start as number) + // conversation (12 chars) should be 6x the duration of hi (2 chars) + expect(d1 / d0).toBeCloseTo(6, 0) + }) + + it('applies padding when space allows', () => { + // With enough space, the first filled word should not start exactly at the left anchor + const t = transcript([ + { + start: 0, + end: 100, + text: 'a xxxx here ', + words: [w('a', 0, 5, 0.9), w('xxxx'), w('here', 95, 100, 0.8)], + }, + ]) + + fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const filled = words[1] + // Should be padded away from the left anchor (5) + expect(filled.start).toBeGreaterThan(5) + // And padded away from the right anchor (95) + expect(filled.end).toBeLessThan(95) + }) + + it('skips padding when interval is tight', () => { + // Very short interval between anchors — padding should be skipped + const t = transcript([ + { + start: 0, + end: 1, + text: 'a b c', + words: [w('a', 0, 0.4, 0.9), w('b'), w('c', 0.6, 1, 0.8)], + }, + ]) + + fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const filled = words[1] + // Interval is 0.6-0.4=0.2, charRate=1/5=0.2, padding would be 0.2 which is >= half of 0.2 + // So padding should be skipped and word should use the full interval + expect(filled.start).toBeDefined() + expect(filled.end).toBeDefined() + expect(filled.score).toBeNull() + }) + + it('handles zero-duration gap where anchors are the same', () => { + const t = transcript([ + { + start: 0, + end: 3, + text: 'hello world bye', + words: [w('hello', 0, 1.5, 0.9), w('world'), w('bye', 1.5, 3, 0.8)], + }, + ]) + + fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const filled = words[1] + // Left anchor = 1.5, right anchor = 1.5 — zero interval + expect(filled.start).toBe(1.5) + expect(filled.end).toBe(1.5) + expect(filled.score).toBeNull() + }) + + it('sets score to null for filled words and leaves existing scores untouched', () => { + const t = transcript([ + { + start: 0, + end: 5, + text: 'aaa bbb ccc', + words: [w('aaa', 0, 1, 0.95), w('bbb'), w('ccc', 4, 5, 0.87)], + }, + ]) + + fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + expect(words[0].score).toBe(0.95) + expect(words[1].score).toBeNull() + expect(words[2].score).toBe(0.87) + }) + + it('skips segments without words array gracefully', () => { + const t = transcript([ + { + start: 0, + end: 3, + text: 'hello world', + // no words array + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + expect(stats.totalSegments).toBe(1) + expect(stats.totalWords).toBe(0) + expect(stats.wordsFilled).toBe(0) + }) + + it('returns zero stats for empty transcript', () => { + const t = transcript([]) + + const stats = fillTimingGaps(t, enabledConfig) + + expect(stats).toEqual({ + totalSegments: 0, + totalWords: 0, + wordsFilled: 0, + gapsFound: 0, + gapTypes: { middle: 0, start: 0, end: 0, entireSegment: 0 }, + }) + }) + + it('handles real-world pattern: numbers and currency words', () => { + // Simulates WhisperX output where numbers/currency lack timing + const t = transcript([ + { + start: 10, + end: 16, + text: 'it costs $8 or $0.25 per unit', + words: [ + w('it', 10, 10.5, 0.9), + w('costs', 10.5, 11, 0.85), + w('$8'), // missing timing (currency) + w('or', 12, 12.5, 0.8), + w('$0.25'), // missing timing (currency) + w('per', 14, 14.5, 0.9), + w('unit', 14.5, 16, 0.85), + ], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + expect(stats.wordsFilled).toBe(2) + expect(stats.gapsFound).toBe(2) + + const words = getWords(t) + // $8 should be between 'costs' end (11) and 'or' start (12) + expect(words[2].start).toBeGreaterThanOrEqual(11) + expect(words[2].end).toBeLessThanOrEqual(12) + expect(words[2].score).toBeNull() + + // $0.25 should be between 'or' end (12.5) and 'per' start (14) + expect(words[4].start).toBeGreaterThanOrEqual(12.5) + expect(words[4].end).toBeLessThanOrEqual(14) + expect(words[4].score).toBeNull() + }) + + it('returns early when config is disabled', () => { + const t = transcript([ + { + start: 0, + end: 3, + text: 'hello world', + words: [w('hello'), w('world')], + }, + ]) + + const stats = fillTimingGaps(t, { enabled: false }) + + expect(stats.totalSegments).toBe(0) + expect(stats.wordsFilled).toBe(0) + // Words should be unchanged + const words = getWords(t) + expect(words[0].start).toBeUndefined() + }) + + it('handles partial gap: word with start but missing end', () => { + const t = transcript([ + { + start: 0, + end: 5, + text: 'hello world', + words: [ + w('hello', 0, 1, 0.9), + { word: 'world', start: 2, score: 0.7 } as WhisperxWord, + ], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const filled = words[1] + expect(filled.end).toBeDefined() + expect(filled.end).toBeLessThanOrEqual(5) + expect(filled.score).toBeNull() + expect(stats.wordsFilled).toBe(1) + }) + + it('handles partial gap: word with end but missing start', () => { + const t = transcript([ + { + start: 0, + end: 5, + text: 'hello world', + words: [ + w('hello', 0, 1, 0.9), + { word: 'world', end: 4, score: 0.7 } as WhisperxWord, + ], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + const filled = words[1] + expect(filled.start).toBeDefined() + expect(filled.start).toBeGreaterThanOrEqual(0) + expect(filled.score).toBeNull() + expect(stats.wordsFilled).toBe(1) + }) + + it('clamps partial gap (missing end) to next word start to avoid overlap', () => { + // charRate for "ab cd ef" over 0-10 = 10/8 = 1.25 per char + // estimated end for "cd" = 3 + 1.25*2 = 5.5, but next word starts at 4 + // should clamp to 4 + const t = transcript([ + { + start: 0, + end: 10, + text: 'ab cd ef', + words: [ + w('ab', 0, 2, 0.9), + { word: 'cd', start: 3 } as WhisperxWord, + w('ef', 4, 10, 0.8), + ], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + expect(stats.wordsFilled).toBe(1) + expect(words[1].end).toBeDefined() + // Must not exceed the next word's start (4) + expect(words[1].end as number).toBeLessThanOrEqual(4) + }) + + it('clamps partial gap (missing start) to previous word end to avoid overlap', () => { + // charRate for "ab cd ef" over 0-10 = 10/8 = 1.25 per char + // estimated start for "cd" = 7 - 1.25*2 = 4.5, but previous word ends at 6 + // should clamp to 6 + const t = transcript([ + { + start: 0, + end: 10, + text: 'ab cd ef', + words: [ + w('ab', 0, 6, 0.9), + { word: 'cd', end: 7 } as WhisperxWord, + w('ef', 8, 10, 0.8), + ], + }, + ]) + + const stats = fillTimingGaps(t, enabledConfig) + + const words = getWords(t) + expect(stats.wordsFilled).toBe(1) + expect(words[1].start).toBeDefined() + // Must not precede the previous word's end (6) + expect(words[1].start as number).toBeGreaterThanOrEqual(6) + }) +}) diff --git a/lambdas/pipeline/steps/fill-timing-gaps.ts b/lambdas/pipeline/steps/fill-timing-gaps.ts new file mode 100644 index 0000000..c9dd7b1 --- /dev/null +++ b/lambdas/pipeline/steps/fill-timing-gaps.ts @@ -0,0 +1,260 @@ +import type { FillTimingGapsConfig } from '@podwhisperer/config' +import type { WhisperxResult, WhisperxSegment, WhisperxWord } from '../types' + +/** A gap of consecutive words missing both start and end timestamps */ +interface Gap { + startIdx: number + endIdx: number +} + +/** Classification of where a gap occurs within a segment */ +type GapType = 'start' | 'end' | 'middle' | 'entireSegment' + +/** Stats returned by the fill-timing-gaps step */ +export interface FillTimingGapsStats { + totalSegments: number + totalWords: number + wordsFilled: number + gapsFound: number + gapTypes: { + middle: number + start: number + end: number + entireSegment: number + } +} + +/** + * Finds groups of consecutive words where both start and end are undefined. + */ +function findGaps(words: WhisperxWord[]): Gap[] { + const gaps: Gap[] = [] + let i = 0 + while (i < words.length) { + if (words[i].start === undefined && words[i].end === undefined) { + const startIdx = i + while ( + i < words.length && + words[i].start === undefined && + words[i].end === undefined + ) { + i++ + } + gaps.push({ startIdx, endIdx: i - 1 }) + } else { + i++ + } + } + return gaps +} + +/** + * Classifies a gap based on its position within the words array. + */ +function classifyGap(gap: Gap, totalWords: number): GapType { + if (gap.startIdx === 0 && gap.endIdx === totalWords - 1) { + return 'entireSegment' + } + if (gap.startIdx === 0) { + return 'start' + } + if (gap.endIdx === totalWords - 1) { + return 'end' + } + return 'middle' +} + +/** + * Resolves the left and right time anchors for a gap. + */ +function resolveAnchors( + gap: Gap, + words: WhisperxWord[], + segment: WhisperxSegment, +): { leftAnchor: number; rightAnchor: number } { + const leftAnchor = + gap.startIdx > 0 + ? (words[gap.startIdx - 1].end ?? segment.start) + : segment.start + const rightAnchor = + gap.endIdx < words.length - 1 + ? (words[gap.endIdx + 1].start ?? segment.end) + : segment.end + return { leftAnchor, rightAnchor } +} + +/** + * Computes the character rate for a segment (seconds per character). + */ +function computeCharRate(segment: WhisperxSegment): number { + const duration = segment.end - segment.start + const totalChars = segment.text.length + if (totalChars === 0 || duration <= 0) { + return 0 + } + return duration / totalChars +} + +/** + * Distributes timestamps proportionally across gap words based on character length. + * Applies padding from anchors based on the segment's character rate. + */ +function distributeGap( + words: WhisperxWord[], + gap: Gap, + leftAnchor: number, + rightAnchor: number, + charRate: number, +): void { + const gapWords = words.slice(gap.startIdx, gap.endIdx + 1) + const totalChars = gapWords.reduce((sum, w) => sum + w.word.length, 0) + + if (totalChars === 0) { + // Edge case: all empty words, give them the same timestamp + for (let i = gap.startIdx; i <= gap.endIdx; i++) { + words[i].start = round3(leftAnchor) + words[i].end = round3(leftAnchor) + words[i].score = null + } + return + } + + const interval = rightAnchor - leftAnchor + + // Calculate padding: use charRate as padding from each anchor + let padding = charRate + // Skip padding if it would consume more than half the available interval + if (padding * 2 >= interval) { + padding = 0 + } + + const paddedStart = leftAnchor + padding + const paddedEnd = rightAnchor - padding + const paddedInterval = paddedEnd - paddedStart + + if (paddedInterval <= 0) { + // No room — assign all words the same timestamp + for (let i = gap.startIdx; i <= gap.endIdx; i++) { + words[i].start = round3(leftAnchor) + words[i].end = round3(leftAnchor) + words[i].score = null + } + return + } + + let cursor = paddedStart + for (let i = gap.startIdx; i <= gap.endIdx; i++) { + const proportion = words[i].word.length / totalChars + const duration = paddedInterval * proportion + words[i].start = round3(cursor) + words[i].end = round3(cursor + duration) + words[i].score = null + cursor += duration + } +} + +/** + * Handles words that have only one of start/end missing (partial gaps). + */ +function fillPartialGaps( + words: WhisperxWord[], + segment: WhisperxSegment, + charRate: number, +): number { + let filled = 0 + for (let i = 0; i < words.length; i++) { + const word = words[i] + if (word.start !== undefined && word.end === undefined) { + const estimated = word.start + charRate * word.word.length + const nextStart = words[i + 1]?.start + word.end = round3( + Math.min( + estimated, + segment.end, + ...(nextStart !== undefined ? [nextStart] : []), + ), + ) + word.score = null + filled++ + } else if (word.start === undefined && word.end !== undefined) { + const estimated = word.end - charRate * word.word.length + const prevEnd = words[i - 1]?.end + word.start = round3( + Math.max( + estimated, + segment.start, + ...(prevEnd !== undefined ? [prevEnd] : []), + ), + ) + word.score = null + filled++ + } + } + return filled +} + +/** + * Rounds a number to 3 decimal places (millisecond precision). + */ +function round3(n: number): number { + return Math.round(n * 1000) / 1000 +} + +/** + * Fills missing word timestamps in a WhisperX transcript using + * character-proportional interpolation between anchor points. + * + * Mutates the transcript in place. + * + * @param transcript - WhisperX result to modify in place + * @param config - Fill timing gaps configuration + * @returns Stats about gaps found and filled + */ +export function fillTimingGaps( + transcript: WhisperxResult, + config: FillTimingGapsConfig, +): FillTimingGapsStats { + const stats: FillTimingGapsStats = { + totalSegments: 0, + totalWords: 0, + wordsFilled: 0, + gapsFound: 0, + gapTypes: { middle: 0, start: 0, end: 0, entireSegment: 0 }, + } + + if (!config.enabled) { + return stats + } + + stats.totalSegments = transcript.segments.length + + for (const segment of transcript.segments) { + if (!segment.words || segment.words.length === 0) { + continue + } + + const words = segment.words + stats.totalWords += words.length + + const charRate = computeCharRate(segment) + + // Handle partial gaps first (only start or end missing) + stats.wordsFilled += fillPartialGaps(words, segment, charRate) + + // Find and fill complete gaps (both start and end missing) + const gaps = findGaps(words) + stats.gapsFound += gaps.length + + for (const gap of gaps) { + const gapType = classifyGap(gap, words.length) + stats.gapTypes[gapType]++ + + const { leftAnchor, rightAnchor } = resolveAnchors(gap, words, segment) + distributeGap(words, gap, leftAnchor, rightAnchor, charRate) + + stats.wordsFilled += gap.endIdx - gap.startIdx + 1 + } + } + + return stats +} diff --git a/packages/config/index.test.ts b/packages/config/index.test.ts index 160c0aa..672ce3d 100644 --- a/packages/config/index.test.ts +++ b/packages/config/index.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest' import { CaptionsConfigSchema, + FillTimingGapsConfigSchema, LlmRefinementConfigSchema, NormalizationConfigSchema, PipelineConfigSchema, @@ -110,6 +111,22 @@ describe('CaptionsConfigSchema', () => { }) }) +describe('FillTimingGapsConfigSchema', () => { + it('populates defaults for empty object', () => { + const result = FillTimingGapsConfigSchema.parse({}) + + expect(result).toEqual({ + enabled: true, + }) + }) + + it('allows overriding enabled to false', () => { + const result = FillTimingGapsConfigSchema.parse({ enabled: false }) + + expect(result.enabled).toBe(false) + }) +}) + describe('PipelineConfigSchema', () => { it('populates nested defaults for empty object', () => { const result = PipelineConfigSchema.parse({}) @@ -123,6 +140,9 @@ describe('PipelineConfigSchema', () => { skipIfOutputExists: false, hfTokenSsmPath: '/podwhisperer/hf_token', }) + expect(result.fillTimingGaps).toEqual({ + enabled: true, + }) expect(result.replacementRules).toBeUndefined() expect(result.llmRefinement).toBeUndefined() expect(result.normalization).toEqual({ diff --git a/packages/config/index.ts b/packages/config/index.ts index d1325db..916c335 100644 --- a/packages/config/index.ts +++ b/packages/config/index.ts @@ -136,6 +136,25 @@ export const CaptionsConfigSchema = z.object({ */ export type CaptionsConfig = z.infer +/** + * Schema for fill-timing-gaps configuration. + * Controls the step that fills missing word timestamps using character-proportional interpolation. + */ +export const FillTimingGapsConfigSchema = z.object({ + /** Whether the fill-timing-gaps step is applied (default: true) */ + enabled: z.boolean().default(true), +}) + +/** + * Configuration for fill-timing-gaps step. + */ +export type FillTimingGapsConfig = z.infer + +/** Default values for FillTimingGapsConfig */ +const fillTimingGapsDefaults: FillTimingGapsConfig = { + enabled: true, +} + /** * Known Whisper model names. * Using z.union with literals + z.string() allows known models with autocomplete @@ -248,25 +267,31 @@ export const PipelineConfigSchema = z.object({ ...val, })), - /** STEP 1 (OPTIONAL): Replacement rules for text substitution in transcripts */ + /** STEP 1: Fill timing gaps - fills missing word timestamps using character-proportional interpolation */ + fillTimingGaps: FillTimingGapsConfigSchema.optional().transform((val) => ({ + ...fillTimingGapsDefaults, + ...val, + })), + + /** STEP 2 (OPTIONAL): Replacement rules for text substitution in transcripts */ replacementRules: z.array(ReplacementRuleSchema).optional(), - /** STEP 2 (OPTIONAL): LLM-based refinement of the current transcript via Bedrock */ + /** STEP 3 (OPTIONAL): LLM-based refinement of the current transcript via Bedrock */ llmRefinement: LlmRefinementConfigSchema.optional(), - /** STEP 3: Normalization: makes sure the transcript is properly segmented so each segment is readable enough */ + /** STEP 4: Normalization: makes sure the transcript is properly segmented so each segment is readable enough */ normalization: NormalizationConfigSchema.optional().transform((val) => ({ ...normalizationDefaults, ...val, })), - /** STEP 4: Caption generation settings */ + /** STEP 5: Caption generation settings */ captions: CaptionsConfigSchema.optional().transform((val) => ({ ...captionsDefaults, ...val, })), - /** STEP 5: EventBridge notification on pipeline completion */ + /** STEP 6: EventBridge notification on pipeline completion */ notification: NotificationConfigSchema.optional().transform((val) => ({ ...notificationDefaults, ...val,