diff --git a/README.md b/README.md index d37289c..b8895e7 100644 --- a/README.md +++ b/README.md @@ -302,6 +302,9 @@ 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` | **Examples:** @@ -325,6 +328,15 @@ 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. 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 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..92ec4b2 --- /dev/null +++ b/internal/asr/chunker.go @@ -0,0 +1,141 @@ +// SPDX-FileCopyrightText: 2026 Alby Hernández +// SPDX-License-Identifier: Apache-2.0 + +package asr + +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 + // 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 +} + +// 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 { + 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..e0c4273 --- /dev/null +++ b/internal/asr/chunker_test.go @@ -0,0 +1,206 @@ +// SPDX-FileCopyrightText: 2026 Alby Hernández +// SPDX-License-Identifier: Apache-2.0 + +package asr + +import ( + "errors" + "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 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 + 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..daf2cf7 100644 --- a/internal/asr/transcriber.go +++ b/internal/asr/transcriber.go @@ -186,6 +186,9 @@ type Transcriber struct { vocabSize int blankIdx int maxTokensPerStep int + chunkFrames int64 + overlapFrames int64 + longAudio bool mel *MelFilterbank encoder *ort.DynamicAdvancedSession decoderPool chan *decoderWorker @@ -193,10 +196,21 @@ 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. Enabled +// turns on the windowing; when off, audio over the model limit is rejected. +type ChunkConfig struct { + Enabled bool + Seconds int + OverlapSeconds int } // buildSessionOptions returns the ONNX Runtime session options for the @@ -287,6 +301,26 @@ 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 + 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 libPath := os.Getenv("ONNXRUNTIME_LIB") if libPath == "" { @@ -516,9 +550,31 @@ func (t *Transcriber) transcribe(ctx context.Context, audioData []byte, format, } } - tokens, err := t.runInference(ctx, features, onToken) + subsampling := int64(t.config.SubsamplingFactor) + plan, err := planForAudio(int64(len(features)), t.chunkFrames, t.overlapFrames, subsampling, t.longAudio) if err != nil { - return "", fmt.Errorf("inference failed: %w", err) + 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), "longAudio", t.longAudio) + } + + 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 +619,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 +676,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 +760,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..568a09f 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -43,6 +43,14 @@ 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. 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 @@ -75,6 +83,11 @@ func New(cfg Config) (*Server, error) { Provider: provider, DeviceID: cfg.GPUDeviceID, }, + Chunk: asr.ChunkConfig{ + Enabled: cfg.LongAudio, + 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..99aa902 100644 --- a/main.go +++ b/main.go @@ -32,6 +32,9 @@ 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.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,