From c5f436b203df35b66b4f67ca3f9be9e612f5ac40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alby=20Hern=C3=A1ndez?= Date: Wed, 1 Jul 2026 16:05:27 +0100 Subject: [PATCH 1/2] feat(asr): chunk long audio to stay under the encoder frame limit The exported encoder crashes on audio longer than ~400s: its positional encoding table is [1,9999,1024] centred at 5000, so past 5000 encoder frames the relative-position slice goes out of range and self_attn Add fails on a broadcast mismatch (issues #14, #18). Split the mel feature sequence into overlapping windows sized in seconds (-chunk-seconds, default 300; -chunk-overlap-seconds, default 15), run each through the encoder and TDT decoder, and concatenate the tokens. Each window decodes in full so the LSTM state and previous-token feedback stay coherent, but only emits tokens whose timestep falls in its owned region; ownership of the shared overlap is split at its midpoint so adjacent emit ranges tile the timeline with no gaps or duplicated speech. Streaming filters by the same range, so no buffering is needed. Audio under the chunk size takes the single window path unchanged. planChunks and the frame math are pure and table-tested (tiling invariant, deterministic layout, overlap, validation, model-limit rejection). NewTranscriber rejects chunk sizes that would overrun the model limit. Fixes #14, #18. --- README.md | 9 +++ internal/asr/chunker.go | 119 +++++++++++++++++++++++++++ internal/asr/chunker_test.go | 153 +++++++++++++++++++++++++++++++++++ internal/asr/mel.go | 7 ++ internal/asr/transcriber.go | 73 ++++++++++++++--- internal/server/server.go | 10 +++ main.go | 2 + 7 files changed, 363 insertions(+), 10 deletions(-) create mode 100644 internal/asr/chunker.go create mode 100644 internal/asr/chunker_test.go diff --git a/README.md b/README.md index d37289c..d864471 100644 --- a/README.md +++ b/README.md @@ -302,6 +302,8 @@ chunk_%03d.wav`) or use a CPU image, which is bounded by system RAM instead. | `-ffmpeg-timeout` | Maximum wall-clock time for a single ffmpeg conversion | `60s` | `-ffmpeg-timeout 30s` | | `-gpu` | Execution provider: `cpu` or `cuda` | `cpu` | `-gpu cuda` | | `-gpu-device` | GPU device index for `cuda` | `0` | `-gpu-device 1` | +| `-chunk-seconds` | Sliding-window size for long audio, in seconds | `300` | `-chunk-seconds 240` | +| `-chunk-overlap-seconds` | Overlap between consecutive chunks, in seconds | `15` | `-chunk-overlap-seconds 10` | **Examples:** @@ -325,6 +327,13 @@ chunk_%03d.wav`) or use a CPU image, which is bounded by system RAM instead. ./parakeet -log-level debug 2>&1 | grep -v "Schema error" ``` +### Long Audio + +The model's encoder tops out at 400 seconds of audio in a single pass. Parakeet +splits longer input into overlapping windows (`-chunk-seconds`, `-chunk-overlap-seconds`), +transcribes each, and stitches the results, dropping the overlap so words at the +seams are not duplicated. Files under the chunk size are transcribed in one pass. + ### Environment Variables Every command-line flag also reads from an environment variable: take the flag diff --git a/internal/asr/chunker.go b/internal/asr/chunker.go new file mode 100644 index 0000000..ae4fbef --- /dev/null +++ b/internal/asr/chunker.go @@ -0,0 +1,119 @@ +// SPDX-FileCopyrightText: 2026 Alby Hernández +// SPDX-License-Identifier: Apache-2.0 + +package asr + +import "fmt" + +const ( + // DefaultChunkSeconds and DefaultChunkOverlapSeconds are the out-of-the-box + // window and overlap sizes when a caller leaves them unset. + DefaultChunkSeconds = 300 + DefaultChunkOverlapSeconds = 15 + + // modelMaxEncoderFrames is the longest encoder-frame sequence the exported + // ONNX model accepts: its positional-encoding table is [1, 9999, 1024], + // centred at 5000, so beyond 5000 frames the relative-position slice goes + // out of range and the self-attention Add crashes on a broadcast mismatch. + modelMaxEncoderFrames int64 = 5000 + + // chunkSafetyMarginEncoderFrames keeps windows clear of the hard model + // limit so rounding at the subsampling boundary can never reach it. + chunkSafetyMarginEncoderFrames int64 = 200 +) + +// chunkWindow describes one analysis window over the mel feature sequence. +// +// All fields are absolute mel-frame indices into the full utterance, half-open +// [start, end). The encoder is run over [start, end); to avoid duplicating +// speech in the overlap between adjacent windows, only tokens whose position +// falls in [emitStart, emitEnd) are kept. Adjacent emit ranges tile the whole +// timeline with no gaps or overlaps: emitEnd of one window equals emitStart of +// the next. +type chunkWindow struct { + start int64 + end int64 + emitStart int64 + emitEnd int64 +} + +// planChunks splits a mel sequence of total frames into overlapping windows no +// larger than chunkFrames, each sharing overlapFrames with its neighbours. +// +// The overlap gives the encoder acoustic context and the decoder LSTM time to +// warm up before its owned region begins. Ownership of the shared overlap is +// split at its midpoint: the earlier window emits up to the midpoint, the later +// window from the midpoint on, so every frame is emitted by exactly one window. +// +// When the audio fits in a single window (total <= chunkFrames), it returns one +// window covering everything, which reproduces the non-chunked behaviour. +// +// Callers must pass chunkFrames > overlapFrames >= 0 and total > 0. +func planChunks(total, chunkFrames, overlapFrames int64) []chunkWindow { + if total <= chunkFrames { + return []chunkWindow{{start: 0, end: total, emitStart: 0, emitEnd: total}} + } + + stride := chunkFrames - overlapFrames + + var windows []chunkWindow + for start := int64(0); start < total; start += stride { + end := start + chunkFrames + if end >= total { + end = total + } + windows = append(windows, chunkWindow{start: start, end: end}) + if end == total { + break + } + } + + // Assign emit ranges. The first window owns from 0, the last owns to the + // end, and every interior boundary sits at the midpoint of the overlap + // shared by the two windows straddling it. + for i := range windows { + if i == 0 { + windows[i].emitStart = 0 + } else { + windows[i].emitStart = windows[i-1].emitEnd + } + + if i == len(windows)-1 { + windows[i].emitEnd = total + } else { + // Midpoint of the overlap between window i and i+1. + windows[i].emitEnd = (windows[i+1].start + windows[i].end) / 2 + } + } + + return windows +} + +// melToEncoderFrame converts a mel-frame offset to its encoder-frame index under +// the given subsampling factor. The encoder collapses subsampling mel frames +// into one, so the mapping is integer division. +func melToEncoderFrame(melOffset, subsampling int64) int64 { + if subsampling <= 0 { + return melOffset + } + return melOffset / subsampling +} + +// validateChunking rejects window sizes that would break planChunks or overrun +// the model's positional-encoding limit. Sizes are in mel frames. +func validateChunking(chunkFrames, overlapFrames, subsampling int64) error { + if chunkFrames <= 0 { + return fmt.Errorf("chunk size must be positive, got %d frames", chunkFrames) + } + if overlapFrames < 0 { + return fmt.Errorf("chunk overlap must not be negative, got %d frames", overlapFrames) + } + if overlapFrames >= chunkFrames { + return fmt.Errorf("chunk overlap (%d frames) must be smaller than chunk size (%d frames)", overlapFrames, chunkFrames) + } + maxFrames := (modelMaxEncoderFrames - chunkSafetyMarginEncoderFrames) * subsampling + if chunkFrames > maxFrames { + return fmt.Errorf("chunk size (%d frames) exceeds the model-safe maximum (%d frames)", chunkFrames, maxFrames) + } + return nil +} diff --git a/internal/asr/chunker_test.go b/internal/asr/chunker_test.go new file mode 100644 index 0000000..27961a4 --- /dev/null +++ b/internal/asr/chunker_test.go @@ -0,0 +1,153 @@ +// SPDX-FileCopyrightText: 2026 Alby Hernández +// SPDX-License-Identifier: Apache-2.0 + +package asr + +import ( + "reflect" + "testing" +) + +// planChunks must return a single full-coverage window when the audio fits. +func TestPlanChunks_SingleWindowWhenAudioFits(t *testing.T) { + tests := []struct { + name string + total int64 + }{ + {"well under chunk", 1000}, + {"exactly chunk", 4000}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := planChunks(tt.total, 4000, 500) + + want := []chunkWindow{{start: 0, end: tt.total, emitStart: 0, emitEnd: tt.total}} + if !reflect.DeepEqual(got, want) { + t.Fatalf("planChunks(%d) = %+v, want %+v", tt.total, got, want) + } + }) + } +} + +// planChunks must lay out the known windows and emit ranges for a hand-checked +// case. This is the canary: a wrong midpoint or stride changes these numbers. +func TestPlanChunks_DeterministicLayout(t *testing.T) { + got := planChunks(100, 40, 10) + + want := []chunkWindow{ + {start: 0, end: 40, emitStart: 0, emitEnd: 35}, + {start: 30, end: 70, emitStart: 35, emitEnd: 65}, + {start: 60, end: 100, emitStart: 65, emitEnd: 100}, + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("planChunks(100,40,10) = %+v, want %+v", got, want) + } +} + +// planChunks must produce emit ranges that tile [0,total) with no gap or +// overlap, so every frame is transcribed by exactly one window. +func TestPlanChunks_EmitRangesTileTimeline(t *testing.T) { + cases := []struct { + name string + total int64 + chunkFrames int64 + overlapFrames int64 + }{ + {"even multiple", 90, 40, 10}, + {"one over chunk", 4001, 4000, 500}, + {"large overlap", 100, 40, 35}, + {"no overlap", 100, 40, 0}, + {"many windows", 100000, 30000, 1500}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + windows := planChunks(tc.total, tc.chunkFrames, tc.overlapFrames) + + if len(windows) == 0 { + t.Fatal("no windows produced") + } + if windows[0].emitStart != 0 { + t.Errorf("first emitStart = %d, want 0", windows[0].emitStart) + } + if last := windows[len(windows)-1]; last.emitEnd != tc.total { + t.Errorf("last emitEnd = %d, want %d", last.emitEnd, tc.total) + } + for i, w := range windows { + if w.emitStart > w.emitEnd { + t.Errorf("window %d has emitStart %d > emitEnd %d", i, w.emitStart, w.emitEnd) + } + if w.start > w.emitStart || w.emitEnd > w.end { + t.Errorf("window %d emit range [%d,%d) not inside window [%d,%d)", i, w.emitStart, w.emitEnd, w.start, w.end) + } + if w.end-w.start > tc.chunkFrames { + t.Errorf("window %d spans %d frames, exceeds chunk %d", i, w.end-w.start, tc.chunkFrames) + } + if i > 0 && windows[i-1].emitEnd != w.emitStart { + t.Errorf("gap/overlap between window %d emitEnd %d and window %d emitStart %d", i-1, windows[i-1].emitEnd, i, w.emitStart) + } + } + }) + } +} + +// planChunks must share exactly overlapFrames between consecutive windows. +func TestPlanChunks_ConsecutiveWindowsShareOverlap(t *testing.T) { + windows := planChunks(100000, 30000, 1500) + + for i := 1; i < len(windows); i++ { + // Overlap is where the previous window's end meets this one's start. + overlap := windows[i-1].end - windows[i].start + // The final window is truncated at total, so only assert on full ones. + if windows[i].end-windows[i].start == 30000 && overlap != 1500 { + t.Errorf("window %d overlap = %d, want 1500", i, overlap) + } + } +} + +func TestValidateChunking(t *testing.T) { + const subsampling = 8 + tests := []struct { + name string + chunkFrames int64 + overlapFrames int64 + wantErr bool + }{ + {"valid", 30000, 1500, false}, + {"zero overlap valid", 30000, 0, false}, + {"zero chunk rejected", 0, 0, true}, + {"negative overlap rejected", 30000, -1, true}, + {"overlap equals chunk rejected", 30000, 30000, true}, + {"overlap exceeds chunk rejected", 1000, 2000, true}, + {"exceeds model limit rejected", 40000, 1500, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateChunking(tt.chunkFrames, tt.overlapFrames, subsampling) + + if (err != nil) != tt.wantErr { + t.Fatalf("validateChunking(%d,%d) err = %v, wantErr %v", tt.chunkFrames, tt.overlapFrames, err, tt.wantErr) + } + }) + } +} + +func TestMelToEncoderFrame(t *testing.T) { + tests := []struct { + name string + melOffset int64 + subsampling int64 + want int64 + }{ + {"exact multiple", 800, 8, 100}, + {"rounds down", 807, 8, 100}, + {"zero offset", 0, 8, 0}, + {"zero subsampling passes through", 42, 0, 42}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := melToEncoderFrame(tt.melOffset, tt.subsampling); got != tt.want { + t.Fatalf("melToEncoderFrame(%d,%d) = %d, want %d", tt.melOffset, tt.subsampling, got, tt.want) + } + }) + } +} diff --git a/internal/asr/mel.go b/internal/asr/mel.go index 57b73ee..23b0872 100644 --- a/internal/asr/mel.go +++ b/internal/asr/mel.go @@ -85,6 +85,13 @@ func (m *MelFilterbank) createMelFilterbank() [][]float64 { return filterbank } +// FramesPerSecond returns how many mel frames one second of audio yields, set +// by the hop length and sample rate. It ties frame counts to wall-clock time +// so chunk sizes can be configured in seconds. +func (m *MelFilterbank) FramesPerSecond() int { + return m.sampleRate / m.hopLength +} + // Extract computes mel filterbank features from audio samples func (m *MelFilterbank) Extract(samples []float32) [][]float32 { numFrames := (len(samples)-m.winLength)/m.hopLength + 1 diff --git a/internal/asr/transcriber.go b/internal/asr/transcriber.go index 0b7f71b..e59a82e 100644 --- a/internal/asr/transcriber.go +++ b/internal/asr/transcriber.go @@ -186,6 +186,8 @@ type Transcriber struct { vocabSize int blankIdx int maxTokensPerStep int + chunkFrames int64 + overlapFrames int64 mel *MelFilterbank encoder *ort.DynamicAdvancedSession decoderPool chan *decoderWorker @@ -193,10 +195,19 @@ type Transcriber struct { } // Options groups optional knobs passed to NewTranscriber. Zero values keep -// the previous behavior: WAV-only input, no ffmpeg conversion, CPU inference. +// the previous behavior: WAV-only input, no ffmpeg conversion, CPU inference, +// default chunk sizes. type Options struct { FFmpeg FFmpegConfig GPU GPUConfig + Chunk ChunkConfig +} + +// ChunkConfig sets the sliding-window sizes that keep long audio within the +// model's frame limit. Zero values fall back to the package defaults. +type ChunkConfig struct { + Seconds int + OverlapSeconds int } // buildSessionOptions returns the ONNX Runtime session options for the @@ -287,6 +298,23 @@ func NewTranscriber(modelsDir string, workers int, opts Options) (*Transcriber, // Initialize mel filterbank t.mel = NewMelFilterbank(t.config.FeaturesSize, 16000) + // Resolve chunk sizes (seconds to mel frames) and reject anything that + // would overrun the model's frame limit. + chunkSeconds := opts.Chunk.Seconds + if chunkSeconds <= 0 { + chunkSeconds = DefaultChunkSeconds + } + overlapSeconds := opts.Chunk.OverlapSeconds + if overlapSeconds < 0 { + overlapSeconds = DefaultChunkOverlapSeconds + } + fps := int64(t.mel.FramesPerSecond()) + t.chunkFrames = int64(chunkSeconds) * fps + t.overlapFrames = int64(overlapSeconds) * fps + if err := validateChunking(t.chunkFrames, t.overlapFrames, int64(t.config.SubsamplingFactor)); err != nil { + return nil, fmt.Errorf("invalid chunk configuration: %w", err) + } + // Initialize ONNX Runtime libPath := os.Getenv("ONNXRUNTIME_LIB") if libPath == "" { @@ -516,9 +544,25 @@ func (t *Transcriber) transcribe(ctx context.Context, audioData []byte, format, } } - tokens, err := t.runInference(ctx, features, onToken) - if err != nil { - return "", fmt.Errorf("inference failed: %w", err) + subsampling := int64(t.config.SubsamplingFactor) + plan := planChunks(int64(len(features)), t.chunkFrames, t.overlapFrames) + + if DebugMode { + slog.Debug("chunk plan", "windows", len(plan), "melFrames", len(features)) + } + + var tokens []int + for _, win := range plan { + // Emit bounds are the window's owned region expressed in the window's + // local encoder frames, so tdtDecode drops the overlap it does not own. + emitStart := melToEncoderFrame(win.emitStart-win.start, subsampling) + emitEnd := melToEncoderFrame(win.emitEnd-win.start, subsampling) + + windowTokens, err := t.runInference(ctx, features[win.start:win.end], emitStart, emitEnd, onToken) + if err != nil { + return "", fmt.Errorf("inference failed: %w", err) + } + tokens = append(tokens, windowTokens...) } if DebugMode { @@ -563,7 +607,7 @@ func (t *Transcriber) loadAudio(data []byte, format string) ([]float32, error) { return parseWAV(wavData) } -func (t *Transcriber) runInference(ctx context.Context, features [][]float32, onToken func(tok int)) ([]int, error) { +func (t *Transcriber) runInference(ctx context.Context, features [][]float32, emitStart, emitEnd int64, onToken func(tok int)) ([]int, error) { batchSize := int64(1) numFeatures := int64(t.config.FeaturesSize) numFrames := int64(len(features)) @@ -620,10 +664,15 @@ func (t *Transcriber) runInference(ctx context.Context, features [][]float32, on // Decoder tensors (encoderOut) must remain alive during tdtDecode. // The defers above fire after tdtDecode returns, so this is safe. - return t.tdtDecode(ctx, encoderOut, actualEncodedLen, onToken) + return t.tdtDecode(ctx, encoderOut, actualEncodedLen, emitStart, emitEnd, onToken) } -func (t *Transcriber) tdtDecode(ctx context.Context, encoderOut []float32, encodedLen int64, onToken func(tok int)) ([]int, error) { +// tdtDecode greedily decodes the encoder output for one window. It decodes the +// whole window so the LSTM state and previous-token feedback stay coherent, but +// only collects and streams tokens whose timestep falls in [emitStart, emitEnd); +// this drops the overlap region owned by an adjacent window. Pass emitStart=0 +// and emitEnd=encodedLen to keep everything. +func (t *Transcriber) tdtDecode(ctx context.Context, encoderOut []float32, encodedLen, emitStart, emitEnd int64, onToken func(tok int)) ([]int, error) { // Acquire a pre-initialized worker. Honor cancellation so a client that // disconnects while all workers are busy does not leak a goroutine. var w *decoderWorker @@ -699,11 +748,15 @@ func (t *Transcriber) tdtDecode(ctx context.Context, encoderOut []float32, encod // Update LSTM states for next step copy(w.state1In.GetData(), w.state1Out.GetData()) copy(w.state2In.GetData(), w.state2Out.GetData()) - tokens = append(tokens, token) prevToken = token emittedTokens++ - if onToken != nil { - onToken(token) + // Collect and stream only tokens this window owns; the rest belong + // to an adjacent window's overlap and would duplicate speech. + if timestep >= emitStart && timestep < emitEnd { + tokens = append(tokens, token) + if onToken != nil { + onToken(token) + } } } diff --git a/internal/server/server.go b/internal/server/server.go index 0d1063f..6a9f47d 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -43,6 +43,12 @@ type Config struct { // GPUDeviceID selects the GPU device index for GPU providers. GPUDeviceID int + + // ChunkSeconds is the sliding-window size for long audio, in seconds. + // ChunkOverlapSeconds is how much consecutive windows share so words at + // the seams keep their context. + ChunkSeconds int + ChunkOverlapSeconds int } // Server represents the HTTP server for the ASR service @@ -75,6 +81,10 @@ func New(cfg Config) (*Server, error) { Provider: provider, DeviceID: cfg.GPUDeviceID, }, + Chunk: asr.ChunkConfig{ + Seconds: cfg.ChunkSeconds, + OverlapSeconds: cfg.ChunkOverlapSeconds, + }, }) if err != nil { return nil, fmt.Errorf("failed to initialize transcriber: %w", err) diff --git a/main.go b/main.go index 3884331..512b631 100644 --- a/main.go +++ b/main.go @@ -32,6 +32,8 @@ func main() { flag.DurationVar(&cfg.FFmpegTimeout, "ffmpeg-timeout", 60*time.Second, "Maximum wall-clock time for a single ffmpeg conversion") flag.StringVar(&cfg.GPUProvider, "gpu", "cpu", "Execution provider: cpu or cuda") flag.IntVar(&cfg.GPUDeviceID, "gpu-device", 0, "GPU device index for cuda") + flag.IntVar(&cfg.ChunkSeconds, "chunk-seconds", 300, "Sliding-window size in seconds for long audio (must stay under the model limit)") + flag.IntVar(&cfg.ChunkOverlapSeconds, "chunk-overlap-seconds", 15, "Overlap in seconds between consecutive chunks") flag.Parse() // Any flag not set on the command line falls back to its matching env var, From 2fb27c023a54485a5f72bf4f1c78d4e59081b72c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alby=20Hern=C3=A1ndez?= Date: Wed, 1 Jul 2026 16:42:33 +0100 Subject: [PATCH 2/2] feat(asr): gate long-audio chunking behind --long-audio (default off) Chunking is now opt-in. With -long-audio off (the default) audio over the model's single-pass limit is rejected with a clear error and a log warning pointing at the flag, instead of chunking or crashing. With it on, the overlapping-window path runs as before. planForAudio picks the coverage (single window, chunked, or ErrAudioTooLong) and is pure and table-tested. Chunk-size validation only runs when the mode is enabled. --- README.md | 11 +++++--- internal/asr/chunker.go | 24 +++++++++++++++- internal/asr/chunker_test.go | 53 ++++++++++++++++++++++++++++++++++++ internal/asr/transcriber.go | 22 +++++++++++---- internal/server/server.go | 5 +++- main.go | 1 + 6 files changed, 105 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index d864471..b8895e7 100644 --- a/README.md +++ b/README.md @@ -302,6 +302,7 @@ chunk_%03d.wav`) or use a CPU image, which is bounded by system RAM instead. | `-ffmpeg-timeout` | Maximum wall-clock time for a single ffmpeg conversion | `60s` | `-ffmpeg-timeout 30s` | | `-gpu` | Execution provider: `cpu` or `cuda` | `cpu` | `-gpu cuda` | | `-gpu-device` | GPU device index for `cuda` | `0` | `-gpu-device 1` | +| `-long-audio` | Split audio over the model limit into chunks instead of rejecting it | `false` | `-long-audio` | | `-chunk-seconds` | Sliding-window size for long audio, in seconds | `300` | `-chunk-seconds 240` | | `-chunk-overlap-seconds` | Overlap between consecutive chunks, in seconds | `15` | `-chunk-overlap-seconds 10` | @@ -329,10 +330,12 @@ chunk_%03d.wav`) or use a CPU image, which is bounded by system RAM instead. ### Long Audio -The model's encoder tops out at 400 seconds of audio in a single pass. Parakeet -splits longer input into overlapping windows (`-chunk-seconds`, `-chunk-overlap-seconds`), -transcribes each, and stitches the results, dropping the overlap so words at the -seams are not duplicated. Files under the chunk size are transcribed in one pass. +The model's encoder tops out at 400 seconds of audio in a single pass. By +default, longer input is rejected with a clear error (and a log line pointing +here). Pass `-long-audio` to split it into overlapping windows +(`-chunk-seconds`, `-chunk-overlap-seconds`), transcribe each, and stitch the +results, dropping the overlap so words at the seams are not duplicated. Files +under the chunk size are transcribed in one pass either way. ### Environment Variables diff --git a/internal/asr/chunker.go b/internal/asr/chunker.go index ae4fbef..92ec4b2 100644 --- a/internal/asr/chunker.go +++ b/internal/asr/chunker.go @@ -3,7 +3,14 @@ package asr -import "fmt" +import ( + "errors" + "fmt" +) + +// ErrAudioTooLong is returned when long-audio mode is off and the input exceeds +// what the model can process in a single pass. +var ErrAudioTooLong = errors.New("audio too long for a single pass; enable long-audio mode") const ( // DefaultChunkSeconds and DefaultChunkOverlapSeconds are the out-of-the-box @@ -99,6 +106,21 @@ func melToEncoderFrame(melOffset, subsampling int64) int64 { return melOffset / subsampling } +// planForAudio decides how to cover a mel sequence of total frames. With long +// audio enabled it splits into overlapping windows (planChunks). With it off it +// returns a single full-coverage window, or ErrAudioTooLong when the input would +// overrun the model's single-pass limit, so the caller fails cleanly instead of +// letting the encoder crash on an out-of-range positional-encoding slice. +func planForAudio(total, chunkFrames, overlapFrames, subsampling int64, longAudio bool) ([]chunkWindow, error) { + if longAudio { + return planChunks(total, chunkFrames, overlapFrames), nil + } + if melToEncoderFrame(total, subsampling) > modelMaxEncoderFrames { + return nil, ErrAudioTooLong + } + return []chunkWindow{{start: 0, end: total, emitStart: 0, emitEnd: total}}, nil +} + // validateChunking rejects window sizes that would break planChunks or overrun // the model's positional-encoding limit. Sizes are in mel frames. func validateChunking(chunkFrames, overlapFrames, subsampling int64) error { diff --git a/internal/asr/chunker_test.go b/internal/asr/chunker_test.go index 27961a4..e0c4273 100644 --- a/internal/asr/chunker_test.go +++ b/internal/asr/chunker_test.go @@ -4,6 +4,7 @@ package asr import ( + "errors" "reflect" "testing" ) @@ -131,6 +132,58 @@ func TestValidateChunking(t *testing.T) { } } +func TestPlanForAudio(t *testing.T) { + const ( + subsampling = 8 + chunk = 30000 + overlap = 1500 + // modelMaxEncoderFrames is 5000, so 5000*8 = 40000 mel frames is the + // single-pass ceiling. + underLimit = 40000 + overLimit = 40008 + ) + t.Run("long audio off, short audio, single window", func(t *testing.T) { + plan, err := planForAudio(underLimit, chunk, overlap, subsampling, false) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(plan) != 1 || plan[0].start != 0 || plan[0].end != underLimit { + t.Fatalf("want one full window, got %+v", plan) + } + }) + + t.Run("long audio off, long audio, rejected", func(t *testing.T) { + _, err := planForAudio(overLimit, chunk, overlap, subsampling, false) + + if !errors.Is(err, ErrAudioTooLong) { + t.Fatalf("want ErrAudioTooLong, got %v", err) + } + }) + + t.Run("long audio on, long audio, chunked", func(t *testing.T) { + plan, err := planForAudio(overLimit*3, chunk, overlap, subsampling, true) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(plan) < 2 { + t.Fatalf("want multiple windows for long audio, got %d", len(plan)) + } + }) + + t.Run("long audio on, short audio, single window", func(t *testing.T) { + plan, err := planForAudio(1000, chunk, overlap, subsampling, true) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(plan) != 1 { + t.Fatalf("want one window, got %d", len(plan)) + } + }) +} + func TestMelToEncoderFrame(t *testing.T) { tests := []struct { name string diff --git a/internal/asr/transcriber.go b/internal/asr/transcriber.go index e59a82e..daf2cf7 100644 --- a/internal/asr/transcriber.go +++ b/internal/asr/transcriber.go @@ -188,6 +188,7 @@ type Transcriber struct { maxTokensPerStep int chunkFrames int64 overlapFrames int64 + longAudio bool mel *MelFilterbank encoder *ort.DynamicAdvancedSession decoderPool chan *decoderWorker @@ -204,8 +205,10 @@ type Options struct { } // ChunkConfig sets the sliding-window sizes that keep long audio within the -// model's frame limit. Zero values fall back to the package defaults. +// model's frame limit. Zero values fall back to the package defaults. Enabled +// turns on the windowing; when off, audio over the model limit is rejected. type ChunkConfig struct { + Enabled bool Seconds int OverlapSeconds int } @@ -311,8 +314,11 @@ func NewTranscriber(modelsDir string, workers int, opts Options) (*Transcriber, fps := int64(t.mel.FramesPerSecond()) t.chunkFrames = int64(chunkSeconds) * fps t.overlapFrames = int64(overlapSeconds) * fps - if err := validateChunking(t.chunkFrames, t.overlapFrames, int64(t.config.SubsamplingFactor)); err != nil { - return nil, fmt.Errorf("invalid chunk configuration: %w", err) + t.longAudio = opts.Chunk.Enabled + if t.longAudio { + if err := validateChunking(t.chunkFrames, t.overlapFrames, int64(t.config.SubsamplingFactor)); err != nil { + return nil, fmt.Errorf("invalid chunk configuration: %w", err) + } } // Initialize ONNX Runtime @@ -545,10 +551,16 @@ func (t *Transcriber) transcribe(ctx context.Context, audioData []byte, format, } subsampling := int64(t.config.SubsamplingFactor) - plan := planChunks(int64(len(features)), t.chunkFrames, t.overlapFrames) + plan, err := planForAudio(int64(len(features)), t.chunkFrames, t.overlapFrames, subsampling, t.longAudio) + if err != nil { + slog.Warn("audio exceeds the single-pass model limit; enable --long-audio to transcribe long files in overlapping chunks", + "seconds", float64(len(features))/float64(t.mel.FramesPerSecond()), + "limitSeconds", float64(modelMaxEncoderFrames*subsampling)/float64(t.mel.FramesPerSecond())) + return "", err + } if DebugMode { - slog.Debug("chunk plan", "windows", len(plan), "melFrames", len(features)) + slog.Debug("chunk plan", "windows", len(plan), "melFrames", len(features), "longAudio", t.longAudio) } var tokens []int diff --git a/internal/server/server.go b/internal/server/server.go index 6a9f47d..568a09f 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -46,9 +46,11 @@ type Config struct { // ChunkSeconds is the sliding-window size for long audio, in seconds. // ChunkOverlapSeconds is how much consecutive windows share so words at - // the seams keep their context. + // the seams keep their context. LongAudio enables the windowing; when off, + // audio over the model limit is rejected instead of chunked. ChunkSeconds int ChunkOverlapSeconds int + LongAudio bool } // Server represents the HTTP server for the ASR service @@ -82,6 +84,7 @@ func New(cfg Config) (*Server, error) { DeviceID: cfg.GPUDeviceID, }, Chunk: asr.ChunkConfig{ + Enabled: cfg.LongAudio, Seconds: cfg.ChunkSeconds, OverlapSeconds: cfg.ChunkOverlapSeconds, }, diff --git a/main.go b/main.go index 512b631..99aa902 100644 --- a/main.go +++ b/main.go @@ -34,6 +34,7 @@ func main() { flag.IntVar(&cfg.GPUDeviceID, "gpu-device", 0, "GPU device index for cuda") flag.IntVar(&cfg.ChunkSeconds, "chunk-seconds", 300, "Sliding-window size in seconds for long audio (must stay under the model limit)") flag.IntVar(&cfg.ChunkOverlapSeconds, "chunk-overlap-seconds", 15, "Overlap in seconds between consecutive chunks") + flag.BoolVar(&cfg.LongAudio, "long-audio", false, "Split audio longer than the model limit into overlapping chunks instead of rejecting it") flag.Parse() // Any flag not set on the command line falls back to its matching env var,