diff --git a/README.md b/README.md index f9daa59..cc0138b 100644 --- a/README.md +++ b/README.md @@ -35,19 +35,19 @@ The download installs the Parakeet TDT v2 int8 model and Silero VAD into `models Use VAD mode for normal live captions. It shows a temporary `…` caption while speech is active, then prints a timestamped final caption after 0.5 seconds of silence: ```bash -go run ./src --mode=vad +go run ./src/cmd/captioneer --mode=vad ``` Use fixed mode to decode exact time slices, even through silence: ```bash -go run ./src --mode=fixed --chunk-duration=1s +go run ./src/cmd/captioneer --mode=fixed --chunk-duration=1s ``` Build a reusable executable instead of using `go run`: ```bash -go build -o captioneer ./src +go build -o captioneer ./src/cmd/captioneer ./captioneer --mode=vad ``` @@ -70,27 +70,27 @@ Examples: ```bash # Faster fixed chunks -go run ./src --mode=fixed --chunk-duration=500ms +go run ./src/cmd/captioneer --mode=fixed --chunk-duration=500ms # Fewer VAD preview redraws and four CPU threads -go run ./src --mode=vad --chunk-duration=2s --threads=4 +go run ./src/cmd/captioneer --mode=vad --chunk-duration=2s --threads=4 # Store or reuse models outside this repository -go run ./src --mode=vad --models-dir=/path/to/models +go run ./src/cmd/captioneer --mode=vad --models-dir=/path/to/models ``` There are no environment variables to configure. All supported configuration is supplied with command-line options. ## Source layout -All Go code lives in `src/`. It remains one executable package, split into focused files: +All Go code lives in `src/` as focused packages: -- `main.go` coordinates startup and shutdown. -- `config.go` defines and validates command-line options. -- `capture.go` reads the default system-audio monitor. -- `audio.go` contains PCM conversion and audio calculations. -- `models.go` validates and initializes Parakeet and Silero. -- `captions.go` handles fixed chunks, VAD segments, and terminal output. +- `cmd/captioneer` is the executable entrypoint. +- `config` defines command-line options. +- `audio` contains PCM conversion and audio calculations. +- `capture` reads the default system-audio monitor. +- `models` owns Parakeet and Silero model resources. +- `captions` handles fixed chunks, VAD segments, and terminal output. ## Output and troubleshooting @@ -103,7 +103,7 @@ Final captions go to standard output in this form: Startup messages and overload warnings go to standard error, so you can save final captions separately: ```bash -go run ./src --mode=vad > captions.txt +go run ./src/cmd/captioneer --mode=vad > captions.txt ``` - **Missing model files:** run `./scripts/download-models.sh` again. @@ -115,5 +115,5 @@ go run ./src --mode=vad > captions.txt ```bash go test ./... -go build -o captioneer ./src +go build -o captioneer ./src/cmd/captioneer ``` diff --git a/src/audio.go b/src/audio.go deleted file mode 100644 index e1fa135..0000000 --- a/src/audio.go +++ /dev/null @@ -1,51 +0,0 @@ -package main - -import ( - "encoding/binary" - "fmt" - "math" - "time" -) - -const ( - sampleRate = 16000 - channels = 1 - bytesPerSample = 2 - packetDuration = 100 * time.Millisecond -) - -func samplesPerPacket() int { - return durationSamples(packetDuration) -} - -func pcm16LEToFloat32(pcm []byte) []float32 { - samples := make([]float32, len(pcm)/bytesPerSample) - for i := range samples { - samples[i] = float32(int16(binary.LittleEndian.Uint16(pcm[i*2:]))) / 32768 - } - return samples -} - -func durationSamples(duration time.Duration) int { - return int(duration * sampleRate / time.Second) -} - -func rmsDBFS(samples []float32) float64 { - if len(samples) == 0 { - return math.Inf(-1) - } - - var sum float64 - for _, sample := range samples { - sum += float64(sample * sample) - } - rms := math.Sqrt(sum / float64(len(samples))) - if rms == 0 { - return math.Inf(-1) - } - return 20 * math.Log10(rms) -} - -func formatAudioTime(samples int) string { - return fmt.Sprintf("%.2fs", float64(samples)/sampleRate) -} diff --git a/src/audio/audio.go b/src/audio/audio.go new file mode 100644 index 0000000..3294714 --- /dev/null +++ b/src/audio/audio.go @@ -0,0 +1,52 @@ +// Package audio contains the application's PCM and audio-level utilities. +package audio + +import ( + "encoding/binary" + "fmt" + "math" + "time" +) + +const ( + SampleRate = 16000 + Channels = 1 + BytesPerSample = 2 + PacketDuration = 100 * time.Millisecond +) + +func SamplesPerPacket() int { + return DurationSamples(PacketDuration) +} + +func PCM16LEToFloat32(pcm []byte) []float32 { + samples := make([]float32, len(pcm)/BytesPerSample) + for i := range samples { + samples[i] = float32(int16(binary.LittleEndian.Uint16(pcm[i*2:]))) / 32768 + } + return samples +} + +func DurationSamples(duration time.Duration) int { + return int(duration * SampleRate / time.Second) +} + +func RMSDBFS(samples []float32) float64 { + if len(samples) == 0 { + return math.Inf(-1) + } + + var sum float64 + for _, sample := range samples { + sum += float64(sample * sample) + } + rms := math.Sqrt(sum / float64(len(samples))) + if rms == 0 { + return math.Inf(-1) + } + return 20 * math.Log10(rms) +} + +func FormatTime(samples int) string { + return fmt.Sprintf("%.2fs", float64(samples)/SampleRate) +} diff --git a/src/audio_test.go b/src/audio/audio_test.go similarity index 67% rename from src/audio_test.go rename to src/audio/audio_test.go index 7d89719..e062d01 100644 --- a/src/audio_test.go +++ b/src/audio/audio_test.go @@ -1,4 +1,4 @@ -package main +package audio import ( "encoding/binary" @@ -12,7 +12,7 @@ func TestPCM16LEToFloat32(t *testing.T) { binary.LittleEndian.PutUint16(pcm[0:], uint16(0)) binary.LittleEndian.PutUint16(pcm[2:], uint16(32767)) binary.LittleEndian.PutUint16(pcm[4:], uint16(32768)) - got := pcm16LEToFloat32(pcm) + got := PCM16LEToFloat32(pcm) want := []float32{0, 32767.0 / 32768.0, -1} for i := range want { if math.Abs(float64(got[i]-want[i])) > 0.00001 { @@ -22,19 +22,19 @@ func TestPCM16LEToFloat32(t *testing.T) { } func TestDurationSamples(t *testing.T) { - if got := durationSamples(time.Second); got != sampleRate { - t.Fatalf("one second = %d samples, want %d", got, sampleRate) + if got := DurationSamples(time.Second); got != SampleRate { + t.Fatalf("one second = %d samples, want %d", got, SampleRate) } - if got := durationSamples(1500 * time.Millisecond); got != 24000 { + if got := DurationSamples(1500 * time.Millisecond); got != 24000 { t.Fatalf("1.5 seconds = %d samples, want 24000", got) } } func TestRMSDBFS(t *testing.T) { - if got := rmsDBFS([]float32{1, -1}); math.Abs(got) > 0.00001 { + if got := RMSDBFS([]float32{1, -1}); math.Abs(got) > 0.00001 { t.Fatalf("full-scale RMS = %f dBFS, want 0", got) } - if got := rmsDBFS([]float32{0, 0}); !math.IsInf(got, -1) { + if got := RMSDBFS([]float32{0, 0}); !math.IsInf(got, -1) { t.Fatalf("silence = %f dBFS, want -Inf", got) } } diff --git a/src/captions.go b/src/captions.go deleted file mode 100644 index 0fefea8..0000000 --- a/src/captions.go +++ /dev/null @@ -1,129 +0,0 @@ -package main - -import ( - "fmt" - "io" - "strings" - - sherpa "github.com/k2-fsa/sherpa-onnx-go-linux" -) - -type captionProcessor struct { - cfg settings - recognizer *sherpa.OfflineRecognizer - vad *sherpa.VoiceActivityDetector - output io.Writer - - fixedSamples []float32 - previewSamples []float32 - previewActive bool - previewVisible bool - nextPreviewAt int - totalSamples int -} - -func newCaptionProcessor(cfg settings, recognizer *sherpa.OfflineRecognizer, vad *sherpa.VoiceActivityDetector, output io.Writer) *captionProcessor { - return &captionProcessor{ - cfg: cfg, - recognizer: recognizer, - vad: vad, - output: output, - nextPreviewAt: durationSamples(cfg.chunkDuration), - } -} - -func (p *captionProcessor) accept(samples []float32) { - if len(samples) == 0 { - return - } - if p.cfg.mode == "fixed" { - p.acceptFixed(samples) - return - } - p.acceptVAD(samples) -} - -func (p *captionProcessor) acceptFixed(samples []float32) { - p.fixedSamples = append(p.fixedSamples, samples...) - chunkSize := durationSamples(p.cfg.chunkDuration) - for len(p.fixedSamples) >= chunkSize { - p.emitCommitted(p.totalSamples, p.totalSamples+chunkSize, p.fixedSamples[:chunkSize]) - p.fixedSamples = p.fixedSamples[chunkSize:] - p.totalSamples += chunkSize - } -} - -func (p *captionProcessor) acceptVAD(samples []float32) { - p.totalSamples += len(samples) - p.vad.AcceptWaveform(samples) - - if p.previewActive || rmsDBFS(samples) >= p.cfg.previewThresholdDB { - p.previewActive = true - p.previewSamples = append(p.previewSamples, samples...) - if len(p.previewSamples) >= p.nextPreviewAt { - p.emitPreview(p.previewSamples) - p.nextPreviewAt += durationSamples(p.cfg.chunkDuration) - } - } - p.drainVAD() -} - -func (p *captionProcessor) drainVAD() { - for !p.vad.IsEmpty() { - segment := p.vad.Front() - p.vad.Pop() - p.clearPreview() - p.emitCommitted(segment.Start, segment.Start+len(segment.Samples), segment.Samples) - p.previewSamples = nil - p.previewActive = false - p.nextPreviewAt = durationSamples(p.cfg.chunkDuration) - } -} - -func (p *captionProcessor) flush() { - if p.cfg.mode == "fixed" { - if len(p.fixedSamples) > 0 { - p.emitCommitted(p.totalSamples, p.totalSamples+len(p.fixedSamples), p.fixedSamples) - } - return - } - p.vad.Flush() - p.drainVAD() - p.clearPreview() -} - -func (p *captionProcessor) emitPreview(samples []float32) { - text := p.decode(samples) - if text == "" { - return - } - fmt.Fprintf(p.output, "\r\033[2K… %s", text) - p.previewVisible = true -} - -func (p *captionProcessor) emitCommitted(start, end int, samples []float32) { - text := p.decode(samples) - if text == "" { - return - } - p.clearPreview() - fmt.Fprintf(p.output, "[%s-%s] %s\n", formatAudioTime(start), formatAudioTime(end), text) -} - -func (p *captionProcessor) clearPreview() { - if p.previewVisible { - fmt.Fprint(p.output, "\r\033[2K") - p.previewVisible = false - } -} - -func (p *captionProcessor) decode(samples []float32) string { - if len(samples) == 0 { - return "" - } - stream := sherpa.NewOfflineStream(p.recognizer) - defer sherpa.DeleteOfflineStream(stream) - stream.AcceptWaveform(sampleRate, samples) - p.recognizer.Decode(stream) - return strings.TrimSpace(stream.GetResult().Text) -} diff --git a/src/captions/processor.go b/src/captions/processor.go new file mode 100644 index 0000000..86f503b --- /dev/null +++ b/src/captions/processor.go @@ -0,0 +1,120 @@ +// Package captions turns audio packets into fixed or speech-aware captions. +package captions + +import ( + "fmt" + "io" + + "tea.chunkbyte.com/kato/captioneer/src/audio" + "tea.chunkbyte.com/kato/captioneer/src/config" + "tea.chunkbyte.com/kato/captioneer/src/models" +) + +type Processor struct { + settings config.Settings + resources *models.Resources + output io.Writer + + fixedSamples []float32 + previewSamples []float32 + previewActive bool + previewVisible bool + nextPreviewAt int + totalSamples int +} + +func New(settings config.Settings, resources *models.Resources, output io.Writer) *Processor { + return &Processor{ + settings: settings, + resources: resources, + output: output, + nextPreviewAt: audio.DurationSamples(settings.ChunkDuration), + } +} + +func (p *Processor) Accept(samples []float32) { + if len(samples) == 0 { + return + } + if p.settings.Mode == config.ModeFixed { + p.acceptFixed(samples) + return + } + p.acceptVAD(samples) +} + +func (p *Processor) Flush() { + if p.settings.Mode == config.ModeFixed { + if len(p.fixedSamples) > 0 { + p.emitCommitted(p.totalSamples, p.totalSamples+len(p.fixedSamples), p.fixedSamples) + } + return + } + p.resources.FlushVAD() + p.drainVAD() + p.clearPreview() +} + +func (p *Processor) acceptFixed(samples []float32) { + p.fixedSamples = append(p.fixedSamples, samples...) + chunkSize := audio.DurationSamples(p.settings.ChunkDuration) + for len(p.fixedSamples) >= chunkSize { + p.emitCommitted(p.totalSamples, p.totalSamples+chunkSize, p.fixedSamples[:chunkSize]) + p.fixedSamples = p.fixedSamples[chunkSize:] + p.totalSamples += chunkSize + } +} + +func (p *Processor) acceptVAD(samples []float32) { + p.totalSamples += len(samples) + p.resources.AcceptVAD(samples) + + if p.previewActive || audio.RMSDBFS(samples) >= p.settings.PreviewThresholdDB { + p.previewActive = true + p.previewSamples = append(p.previewSamples, samples...) + if len(p.previewSamples) >= p.nextPreviewAt { + p.emitPreview(p.previewSamples) + p.nextPreviewAt += audio.DurationSamples(p.settings.ChunkDuration) + } + } + p.drainVAD() +} + +func (p *Processor) drainVAD() { + for { + segment, ok := p.resources.NextSpeechSegment() + if !ok { + return + } + p.clearPreview() + p.emitCommitted(segment.Start, segment.Start+len(segment.Samples), segment.Samples) + p.previewSamples = nil + p.previewActive = false + p.nextPreviewAt = audio.DurationSamples(p.settings.ChunkDuration) + } +} + +func (p *Processor) emitPreview(samples []float32) { + text := p.resources.Decode(samples) + if text == "" { + return + } + fmt.Fprintf(p.output, "\r\033[2K… %s", text) + p.previewVisible = true +} + +func (p *Processor) emitCommitted(start, end int, samples []float32) { + text := p.resources.Decode(samples) + if text == "" { + return + } + p.clearPreview() + fmt.Fprintf(p.output, "[%s-%s] %s\n", audio.FormatTime(start), audio.FormatTime(end), text) +} + +func (p *Processor) clearPreview() { + if p.previewVisible { + fmt.Fprint(p.output, "\r\033[2K") + p.previewVisible = false + } +} diff --git a/src/capture.go b/src/capture/pulse.go similarity index 74% rename from src/capture.go rename to src/capture/pulse.go index beb66bc..1a6b809 100644 --- a/src/capture.go +++ b/src/capture/pulse.go @@ -1,4 +1,5 @@ -package main +// Package capture reads the default PulseAudio-compatible output monitor. +package capture import ( "context" @@ -9,9 +10,11 @@ import ( "os/exec" "strings" "sync" + + "tea.chunkbyte.com/kato/captioneer/src/audio" ) -func validatePrograms() error { +func ValidatePrograms() error { for _, program := range []string{"pactl", "parec"} { if _, err := exec.LookPath(program); err != nil { return fmt.Errorf("%s was not found in PATH: %w", program, err) @@ -20,7 +23,7 @@ func validatePrograms() error { return nil } -func defaultMonitorSource(ctx context.Context) (string, error) { +func DefaultMonitorSource(ctx context.Context) (string, error) { output, err := exec.CommandContext(ctx, "pactl", "get-default-sink").Output() if err != nil { return "", err @@ -28,12 +31,12 @@ func defaultMonitorSource(ctx context.Context) (string, error) { return strings.TrimSpace(string(output)) + ".monitor", nil } -func capturePackets(ctx context.Context, monitorSource string) (<-chan []float32, func() error, error) { +func Packets(ctx context.Context, monitorSource string) (<-chan []float32, func() error, error) { cmd := exec.CommandContext(ctx, "parec", "--device="+monitorSource, "--format=s16le", - fmt.Sprintf("--rate=%d", sampleRate), - fmt.Sprintf("--channels=%d", channels), + fmt.Sprintf("--rate=%d", audio.SampleRate), + fmt.Sprintf("--channels=%d", audio.Channels), "--raw", "--latency-msec=50", ) @@ -56,11 +59,11 @@ func capturePackets(ctx context.Context, monitorSource string) (<-chan []float32 go func() { defer close(packets) - packetBytes := make([]byte, samplesPerPacket()*bytesPerSample) + packetBytes := make([]byte, audio.SamplesPerPacket()*audio.BytesPerSample) for { n, readErr := io.ReadFull(stdout, packetBytes) if n > 0 { - enqueueLatest(packets, pcm16LEToFloat32(packetBytes[:n])) + enqueueLatest(packets, audio.PCM16LEToFloat32(packetBytes[:n])) } if readErr != nil { if ctx.Err() == nil && !errors.Is(readErr, io.EOF) && !errors.Is(readErr, io.ErrUnexpectedEOF) { @@ -80,7 +83,6 @@ func enqueueLatest(packets chan []float32, packet []float32) { default: } - // Keep the newest audio so captions recover quickly after overload. select { case <-packets: default: diff --git a/src/main.go b/src/cmd/captioneer/main.go similarity index 50% rename from src/main.go rename to src/cmd/captioneer/main.go index 914c44b..bc163ff 100644 --- a/src/main.go +++ b/src/cmd/captioneer/main.go @@ -8,51 +8,47 @@ import ( "os/signal" "syscall" - sherpa "github.com/k2-fsa/sherpa-onnx-go-linux" + "tea.chunkbyte.com/kato/captioneer/src/captions" + "tea.chunkbyte.com/kato/captioneer/src/capture" + "tea.chunkbyte.com/kato/captioneer/src/config" + "tea.chunkbyte.com/kato/captioneer/src/models" ) func main() { - cfg := parseFlags() - if err := validateSettings(cfg); err != nil { + settings := config.Parse() + if err := settings.Validate(); err != nil { log.Fatal(err) } - if err := validatePrograms(); err != nil { + if err := capture.ValidatePrograms(); err != nil { log.Fatal(err) } - paths, err := validateModels(cfg.modelsDir, cfg.mode) + + resources, err := models.New(settings) if err != nil { log.Fatal(err) } + defer resources.Close() ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - recognizer := newRecognizer(paths, cfg.threads) - defer sherpa.DeleteOfflineRecognizer(recognizer) - - var vad *sherpa.VoiceActivityDetector - if cfg.mode == "vad" { - vad = newVAD(paths.vad, cfg.threads) - defer sherpa.DeleteVoiceActivityDetector(vad) - } - - monitorSource, err := defaultMonitorSource(ctx) + monitorSource, err := capture.DefaultMonitorSource(ctx) if err != nil { log.Fatalf("find default output monitor: %v", err) } fmt.Fprintf(os.Stderr, "Capturing from: %s\n", monitorSource) fmt.Fprintln(os.Stderr, "Press Ctrl+C to stop.") - packets, waitCapture, err := capturePackets(ctx, monitorSource) + packets, waitCapture, err := capture.Packets(ctx, monitorSource) if err != nil { log.Fatalf("start system-audio capture: %v", err) } - processor := newCaptionProcessor(cfg, recognizer, vad, os.Stdout) + processor := captions.New(settings, resources, os.Stdout) for packet := range packets { - processor.accept(packet) + processor.Accept(packet) } - processor.flush() + processor.Flush() if err := waitCapture(); err != nil && ctx.Err() == nil { log.Fatalf("audio capture stopped unexpectedly: %v", err) } diff --git a/src/config.go b/src/config.go deleted file mode 100644 index b9b3939..0000000 --- a/src/config.go +++ /dev/null @@ -1,41 +0,0 @@ -package main - -import ( - "errors" - "flag" - "time" -) - -const defaultModelsDir = "models" - -type settings struct { - mode string - chunkDuration time.Duration - modelsDir string - threads int - previewThresholdDB float64 -} - -func parseFlags() settings { - var cfg settings - flag.StringVar(&cfg.mode, "mode", "", "caption mode: vad or fixed (required)") - flag.DurationVar(&cfg.chunkDuration, "chunk-duration", time.Second, "fixed chunk size or VAD preview refresh interval") - flag.StringVar(&cfg.modelsDir, "models-dir", defaultModelsDir, "directory containing downloaded Sherpa models") - flag.IntVar(&cfg.threads, "threads", 2, "CPU threads used by recognition and VAD") - flag.Float64Var(&cfg.previewThresholdDB, "preview-threshold-dbfs", -45, "RMS dBFS threshold used to begin VAD previews") - flag.Parse() - return cfg -} - -func validateSettings(cfg settings) error { - if cfg.mode != "fixed" && cfg.mode != "vad" { - return errors.New("--mode is required and must be either fixed or vad") - } - if cfg.chunkDuration <= 0 { - return errors.New("--chunk-duration must be positive") - } - if cfg.threads <= 0 { - return errors.New("--threads must be positive") - } - return nil -} diff --git a/src/config/config.go b/src/config/config.go new file mode 100644 index 0000000..f986aab --- /dev/null +++ b/src/config/config.go @@ -0,0 +1,51 @@ +// Package config defines Captioneer's command-line configuration. +package config + +import ( + "errors" + "flag" + "time" +) + +const DefaultModelsDir = "models" + +type Mode string + +const ( + ModeFixed Mode = "fixed" + ModeVAD Mode = "vad" +) + +type Settings struct { + Mode Mode + ChunkDuration time.Duration + ModelsDir string + Threads int + PreviewThresholdDB float64 +} + +func Parse() Settings { + var settings Settings + mode := "" + flag.StringVar(&mode, "mode", "", "caption mode: vad or fixed (required)") + flag.DurationVar(&settings.ChunkDuration, "chunk-duration", time.Second, "fixed chunk size or VAD preview refresh interval") + flag.StringVar(&settings.ModelsDir, "models-dir", DefaultModelsDir, "directory containing downloaded Sherpa models") + flag.IntVar(&settings.Threads, "threads", 2, "CPU threads used by recognition and VAD") + flag.Float64Var(&settings.PreviewThresholdDB, "preview-threshold-dbfs", -45, "RMS dBFS threshold used to begin VAD previews") + flag.Parse() + settings.Mode = Mode(mode) + return settings +} + +func (s Settings) Validate() error { + if s.Mode != ModeFixed && s.Mode != ModeVAD { + return errors.New("--mode is required and must be either fixed or vad") + } + if s.ChunkDuration <= 0 { + return errors.New("--chunk-duration must be positive") + } + if s.Threads <= 0 { + return errors.New("--threads must be positive") + } + return nil +} diff --git a/src/config/config_test.go b/src/config/config_test.go new file mode 100644 index 0000000..b1ead89 --- /dev/null +++ b/src/config/config_test.go @@ -0,0 +1,16 @@ +package config + +import ( + "testing" + "time" +) + +func TestSettingsValidate(t *testing.T) { + valid := Settings{Mode: ModeFixed, ChunkDuration: time.Second, Threads: 1} + if err := valid.Validate(); err != nil { + t.Fatalf("valid settings failed: %v", err) + } + if err := (Settings{Mode: "", ChunkDuration: time.Second, Threads: 1}).Validate(); err == nil { + t.Fatal("missing mode unexpectedly passed") + } +} diff --git a/src/config_test.go b/src/config_test.go deleted file mode 100644 index 351c699..0000000 --- a/src/config_test.go +++ /dev/null @@ -1,26 +0,0 @@ -package main - -import ( - "strings" - "testing" - "time" -) - -func TestValidateSettings(t *testing.T) { - if err := validateSettings(settings{mode: "fixed", chunkDuration: time.Second, threads: 1}); err != nil { - t.Fatalf("valid settings failed: %v", err) - } - if err := validateSettings(settings{mode: "", chunkDuration: time.Second, threads: 1}); err == nil { - t.Fatal("missing mode unexpectedly passed") - } -} - -func TestValidateModelsExplainsHowToInstallThem(t *testing.T) { - _, err := validateModels(t.TempDir(), "vad") - if err == nil { - t.Fatal("missing models unexpectedly passed validation") - } - if !strings.Contains(err.Error(), "./scripts/download-models.sh") { - t.Fatalf("missing-model error did not provide setup command: %v", err) - } -} diff --git a/src/models.go b/src/models.go deleted file mode 100644 index a76ca74..0000000 --- a/src/models.go +++ /dev/null @@ -1,89 +0,0 @@ -package main - -import ( - "fmt" - "log" - "os" - "path/filepath" - "strings" - - sherpa "github.com/k2-fsa/sherpa-onnx-go-linux" -) - -type modelPaths struct { - encoder string - decoder string - joiner string - tokens string - vad string -} - -func validateModels(modelsDir, mode string) (modelPaths, error) { - parakeetDir := filepath.Join(modelsDir, "parakeet-tdt-v2") - paths := modelPaths{ - encoder: filepath.Join(parakeetDir, "encoder.int8.onnx"), - decoder: filepath.Join(parakeetDir, "decoder.int8.onnx"), - joiner: filepath.Join(parakeetDir, "joiner.int8.onnx"), - tokens: filepath.Join(parakeetDir, "tokens.txt"), - vad: filepath.Join(modelsDir, "silero_vad_v5.onnx"), - } - - required := []string{paths.encoder, paths.decoder, paths.joiner, paths.tokens} - if mode == "vad" { - required = append(required, paths.vad) - } - - var missing []string - for _, path := range required { - info, err := os.Stat(path) - if err != nil || info.IsDir() { - missing = append(missing, path) - } - } - if len(missing) > 0 { - return modelPaths{}, fmt.Errorf("missing required model files:\n %s\nRun ./scripts/download-models.sh", strings.Join(missing, "\n ")) - } - return paths, nil -} - -func newRecognizer(paths modelPaths, threads int) *sherpa.OfflineRecognizer { - recognizer := sherpa.NewOfflineRecognizer(&sherpa.OfflineRecognizerConfig{ - FeatConfig: sherpa.FeatureConfig{SampleRate: sampleRate, FeatureDim: 80}, - ModelConfig: sherpa.OfflineModelConfig{ - Transducer: sherpa.OfflineTransducerModelConfig{ - Encoder: paths.encoder, - Decoder: paths.decoder, - Joiner: paths.joiner, - }, - Tokens: paths.tokens, - NumThreads: threads, - Provider: "cpu", - ModelType: "nemo_transducer", - }, - DecodingMethod: "greedy_search", - }) - if recognizer == nil { - log.Fatal("create Parakeet recognizer: Sherpa returned nil") - } - return recognizer -} - -func newVAD(model string, threads int) *sherpa.VoiceActivityDetector { - vad := sherpa.NewVoiceActivityDetector(&sherpa.VadModelConfig{ - SileroVad: sherpa.SileroVadModelConfig{ - Model: model, - Threshold: 0.5, - MinSilenceDuration: 0.5, - MinSpeechDuration: 0.25, - WindowSize: 512, - MaxSpeechDuration: 30, - }, - SampleRate: sampleRate, - NumThreads: threads, - Provider: "cpu", - }, 60) - if vad == nil { - log.Fatal("create Silero VAD: Sherpa returned nil") - } - return vad -} diff --git a/src/models/sherpa.go b/src/models/sherpa.go new file mode 100644 index 0000000..90732ba --- /dev/null +++ b/src/models/sherpa.go @@ -0,0 +1,149 @@ +// Package models owns Sherpa-Onnx model setup and inference resources. +package models + +import ( + "fmt" + "os" + "path/filepath" + "strings" + + sherpa "github.com/k2-fsa/sherpa-onnx-go-linux" + "tea.chunkbyte.com/kato/captioneer/src/audio" + "tea.chunkbyte.com/kato/captioneer/src/config" +) + +type SpeechSegment struct { + Start int + Samples []float32 +} + +type modelPaths struct { + encoder string + decoder string + joiner string + tokens string + vad string +} + +type Resources struct { + recognizer *sherpa.OfflineRecognizer + vad *sherpa.VoiceActivityDetector +} + +func New(settings config.Settings) (*Resources, error) { + paths, err := validatePaths(settings.ModelsDir, settings.Mode) + if err != nil { + return nil, err + } + + recognizer := sherpa.NewOfflineRecognizer(&sherpa.OfflineRecognizerConfig{ + FeatConfig: sherpa.FeatureConfig{SampleRate: audio.SampleRate, FeatureDim: 80}, + ModelConfig: sherpa.OfflineModelConfig{ + Transducer: sherpa.OfflineTransducerModelConfig{ + Encoder: paths.encoder, + Decoder: paths.decoder, + Joiner: paths.joiner, + }, + Tokens: paths.tokens, + NumThreads: settings.Threads, + Provider: "cpu", + ModelType: "nemo_transducer", + }, + DecodingMethod: "greedy_search", + }) + if recognizer == nil { + return nil, fmt.Errorf("create Parakeet recognizer: Sherpa returned nil") + } + + resources := &Resources{recognizer: recognizer} + if settings.Mode == config.ModeVAD { + resources.vad = newVAD(paths.vad, settings.Threads) + if resources.vad == nil { + resources.Close() + return nil, fmt.Errorf("create Silero VAD: Sherpa returned nil") + } + } + return resources, nil +} + +func (r *Resources) Close() { + if r.vad != nil { + sherpa.DeleteVoiceActivityDetector(r.vad) + r.vad = nil + } + if r.recognizer != nil { + sherpa.DeleteOfflineRecognizer(r.recognizer) + r.recognizer = nil + } +} + +func (r *Resources) Decode(samples []float32) string { + if len(samples) == 0 { + return "" + } + stream := sherpa.NewOfflineStream(r.recognizer) + defer sherpa.DeleteOfflineStream(stream) + stream.AcceptWaveform(audio.SampleRate, samples) + r.recognizer.Decode(stream) + return strings.TrimSpace(stream.GetResult().Text) +} + +func (r *Resources) AcceptVAD(samples []float32) { + r.vad.AcceptWaveform(samples) +} + +func (r *Resources) FlushVAD() { + r.vad.Flush() +} + +func (r *Resources) NextSpeechSegment() (SpeechSegment, bool) { + if r.vad.IsEmpty() { + return SpeechSegment{}, false + } + segment := r.vad.Front() + r.vad.Pop() + return SpeechSegment{Start: segment.Start, Samples: segment.Samples}, true +} + +func validatePaths(modelsDir string, mode config.Mode) (modelPaths, error) { + parakeetDir := filepath.Join(modelsDir, "parakeet-tdt-v2") + paths := modelPaths{ + encoder: filepath.Join(parakeetDir, "encoder.int8.onnx"), + decoder: filepath.Join(parakeetDir, "decoder.int8.onnx"), + joiner: filepath.Join(parakeetDir, "joiner.int8.onnx"), + tokens: filepath.Join(parakeetDir, "tokens.txt"), + vad: filepath.Join(modelsDir, "silero_vad_v5.onnx"), + } + required := []string{paths.encoder, paths.decoder, paths.joiner, paths.tokens} + if mode == config.ModeVAD { + required = append(required, paths.vad) + } + + var missing []string + for _, path := range required { + info, err := os.Stat(path) + if err != nil || info.IsDir() { + missing = append(missing, path) + } + } + if len(missing) > 0 { + return modelPaths{}, fmt.Errorf("missing required model files:\n %s\nRun ./scripts/download-models.sh", strings.Join(missing, "\n ")) + } + return paths, nil +} + +func newVAD(model string, threads int) *sherpa.VoiceActivityDetector { + return sherpa.NewVoiceActivityDetector(&sherpa.VadModelConfig{ + SileroVad: sherpa.SileroVadModelConfig{ + Model: model, + Threshold: 0.5, + MinSilenceDuration: 0.5, + MinSpeechDuration: 0.25, + WindowSize: 512, + MaxSpeechDuration: 30, + }, + SampleRate: audio.SampleRate, + NumThreads: threads, + Provider: "cpu", + }, 60) +} diff --git a/src/models/sherpa_test.go b/src/models/sherpa_test.go new file mode 100644 index 0000000..9672e2e --- /dev/null +++ b/src/models/sherpa_test.go @@ -0,0 +1,18 @@ +package models + +import ( + "strings" + "testing" + + "tea.chunkbyte.com/kato/captioneer/src/config" +) + +func TestValidatePathsExplainsHowToInstallModels(t *testing.T) { + _, err := validatePaths(t.TempDir(), config.ModeVAD) + if err == nil { + t.Fatal("missing models unexpectedly passed validation") + } + if !strings.Contains(err.Error(), "./scripts/download-models.sh") { + t.Fatalf("missing-model error did not provide setup command: %v", err) + } +}