From 2f16f8039d22c3c963c1fbf359282cf2bddfe915 Mon Sep 17 00:00:00 2001 From: ciricc Date: Sun, 14 Sep 2025 00:20:10 +0300 Subject: [PATCH] refactor(go bindings): add diarization unit tests --- bindings/go/.gitignore | 1 + bindings/go/Makefile | 11 +++ .../go/examples/go-model-download/main.go | 14 +++- bindings/go/params.go | 5 ++ bindings/go/pkg/whisper/context.go | 3 + bindings/go/pkg/whisper/context_test.go | 71 +++++++++++++++++++ bindings/go/pkg/whisper/interface.go | 7 ++ bindings/go/pkg/whisper/params_wrap.go | 1 + bindings/go/pkg/whisper/util_test.go | 6 +- 9 files changed, 114 insertions(+), 5 deletions(-) diff --git a/bindings/go/.gitignore b/bindings/go/.gitignore index 036df1d3b..20a1b6e48 100644 --- a/bindings/go/.gitignore +++ b/bindings/go/.gitignore @@ -1,2 +1,3 @@ build models +samples/a13.wav diff --git a/bindings/go/Makefile b/bindings/go/Makefile index e4436a6a2..c9ab66255 100644 --- a/bindings/go/Makefile +++ b/bindings/go/Makefile @@ -49,6 +49,9 @@ examples: $(EXAMPLES_DIR) model-small: mkdir examples/go-model-download @${BUILD_DIR}/go-model-download -out models ggml-small.en.bin +model-small-tdrz: mkdir examples/go-model-download + @${BUILD_DIR}/go-model-download -out models ggml-small.en-tdrz.bin + $(EXAMPLES_DIR): mkdir whisper modtidy @echo Build example $(notdir $@) ifeq ($(UNAME_S),Darwin) @@ -57,6 +60,14 @@ else @C_INCLUDE_PATH=${INCLUDE_PATH} LIBRARY_PATH=${LIBRARY_PATH} go build ${BUILD_FLAGS} -o ${BUILD_DIR}/$(notdir $@) ./$@ endif +.PHONY: samples +samples: + @echo "Downloading samples..." + @mkdir -p samples + @wget --quiet --show-progress -O samples/a13.mp3 https://upload.wikimedia.org/wikipedia/commons/transcoded/6/6f/Apollo13-wehaveaproblem.ogg/Apollo13-wehaveaproblem.ogg.mp3 + @ffmpeg -loglevel -0 -y -i samples/a13.mp3 -ar 16000 -ac 1 -c:a pcm_s16le -ss 00:00:00 -to 00:00:30 samples/a13.wav + @rm samples/a13.mp3 + mkdir: @echo Mkdir ${BUILD_DIR} @install -d ${BUILD_DIR} diff --git a/bindings/go/examples/go-model-download/main.go b/bindings/go/examples/go-model-download/main.go index 728c6df53..6ecd1a268 100644 --- a/bindings/go/examples/go-model-download/main.go +++ b/bindings/go/examples/go-model-download/main.go @@ -18,9 +18,10 @@ import ( // CONSTANTS const ( - srcUrl = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/" // The location of the models - srcExt = ".bin" // Filename extension - bufSize = 1024 * 64 // Size of the buffer used for downloading the model + srcUrl = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/" // The location of the models + srcUrlTinydiarize = "https://huggingface.co/akashmjn/tinydiarize-whisper.cpp/resolve/main/" + srcExt = ".bin" // Filename extension + bufSize = 1024 * 64 // Size of the buffer used for downloading the model ) var ( @@ -38,6 +39,7 @@ var ( "large-v2", "large-v2-q5_0", "large-v2-q8_0", "large-v3", "large-v3-q5_0", "large-v3-turbo", "large-v3-turbo-q5_0", "large-v3-turbo-q8_0", + "small.en-tdrz", } ) @@ -219,6 +221,12 @@ func URLForModel(model string) (string, error) { model += srcExt } + srcUrl := srcUrl + + if strings.Contains(model, "tdrz") { + srcUrl = srcUrlTinydiarize + } + // Parse the base URL url, err := url.Parse(srcUrl) if err != nil { diff --git a/bindings/go/params.go b/bindings/go/params.go index 8f669e343..078016493 100644 --- a/bindings/go/params.go +++ b/bindings/go/params.go @@ -47,6 +47,11 @@ func (p *Params) SetPrintTimestamps(v bool) { p.print_timestamps = toBool(v) } +// Enable extra debug information +func (p *Params) SetDebugMode(v bool) { + p.debug_mode = toBool(v) +} + // Enable tinydiarize speaker turn detection func (p *Params) SetDiarize(v bool) { p.tdrz_enable = toBool(v) diff --git a/bindings/go/pkg/whisper/context.go b/bindings/go/pkg/whisper/context.go index 09b35be68..7e2094815 100644 --- a/bindings/go/pkg/whisper/context.go +++ b/bindings/go/pkg/whisper/context.go @@ -3,6 +3,7 @@ package whisper import ( "fmt" "io" + "log" "runtime" "strings" "time" @@ -137,6 +138,8 @@ func (context *context) Process( return fmt.Errorf("lowLevelParams is nil: %w", ErrInternalAppError) } + log.Println("lowLevelParams", lowLevelParams) + st, err := context.st.UnsafeState() if err != nil { return err diff --git a/bindings/go/pkg/whisper/context_test.go b/bindings/go/pkg/whisper/context_test.go index f94c9139d..3c3df9601 100644 --- a/bindings/go/pkg/whisper/context_test.go +++ b/bindings/go/pkg/whisper/context_test.go @@ -1,6 +1,7 @@ package whisper_test import ( + "io" "os" "sync" "testing" @@ -337,6 +338,76 @@ func TestContext_VAD_And_Diarization_Params_DoNotPanic(t *testing.T) { assert.NoError(err) } +func TestDiarization_TwoSpeakers_Boundaries(t *testing.T) { + fh, err := os.Open(MultiSpeakerSamplePath) + require.NoError(t, err) + defer func() { _ = fh.Close() }() + + dec := wav.NewDecoder(fh) + buf, err := dec.FullPCMBuffer() + assert.Equal(t, uint16(1), dec.NumChans) + require.NoError(t, err) + data := buf.AsFloat32Buffer().Data + + model, err := whisper.New(ModelTinydiarizePath) + require.NoError(t, err) + defer func() { _ = model.Close() }() + + // diarize ON with beam search and tighter segmentation + ctxOn, err := model.NewContextWithParams(whisper.SAMPLING_GREEDY, func(p whisper.Parameters) { + p.SetDiarize(true) + p.SetVAD(false) + p.SetSplitOnWord(true) + p.SetMaxSegmentLength(1) + p.SetMaxTokensPerSegment(64) + p.SetTokenTimestamps(true) + }) + require.NoError(t, err) + defer func() { _ = ctxOn.Close() }() + + require.NoError(t, ctxOn.Process(data, nil, nil, nil)) + var turnsOn int + for { + seg, err := ctxOn.NextSegment() + if err == io.EOF { + break + } + require.NoError(t, err) + if seg.SpeakerTurnNext { + turnsOn++ + } + } + require.Greater(t, turnsOn, 0, "expected speaker turn boundaries with diarization enabled") + + // diarize OFF baseline with same segmentation and beam + ctxOff, err := model.NewContextWithParams(whisper.SAMPLING_BEAM_SEARCH, func(p whisper.Parameters) { + p.SetBeamSize(3) + p.SetDiarize(false) + p.SetVAD(false) + p.SetSplitOnWord(true) + p.SetMaxSegmentLength(40) + p.SetMaxTokensPerSegment(64) + p.SetTokenTimestamps(true) + }) + require.NoError(t, err) + defer func() { _ = ctxOff.Close() }() + + require.NoError(t, ctxOff.Process(data, nil, nil, nil)) + var turnsOff int + for { + seg, err := ctxOff.NextSegment() + if err == io.EOF { + break + } + require.NoError(t, err) + if seg.SpeakerTurnNext { + turnsOff++ + } + } + + require.GreaterOrEqual(t, turnsOn, turnsOff, "diarization should not reduce turn boundaries") +} + func TestContext_SpeakerTurnNext_Field_Present(t *testing.T) { assert := assert.New(t) diff --git a/bindings/go/pkg/whisper/interface.go b/bindings/go/pkg/whisper/interface.go index 4aa679c66..9040718e0 100644 --- a/bindings/go/pkg/whisper/interface.go +++ b/bindings/go/pkg/whisper/interface.go @@ -63,6 +63,10 @@ type Model interface { configure ParamsConfigure, ) (Parameters, error) + // Return a new speech-to-text context configured via the provided function + // and sampling strategy. The context is backed by an isolated whisper_state. + NewContextWithParams(sampling SamplingStrategy, configure ParamsConfigure) (Context, error) + // Return true if the model is multilingual. IsMultilingual() bool @@ -107,6 +111,9 @@ type Parameters interface { SetPrintRealtime(bool) SetPrintTimestamps(bool) + // Enable extra debug info (e.g., dump log_mel) + SetDebugMode(bool) + // Diarization (tinydiarize) SetDiarize(bool) diff --git a/bindings/go/pkg/whisper/params_wrap.go b/bindings/go/pkg/whisper/params_wrap.go index ac3a1ccc1..17b70b19e 100644 --- a/bindings/go/pkg/whisper/params_wrap.go +++ b/bindings/go/pkg/whisper/params_wrap.go @@ -41,6 +41,7 @@ func (w *parameters) SetPrintSpecial(v bool) { w.p.SetPrintSpecial(v) func (w *parameters) SetPrintProgress(v bool) { w.p.SetPrintProgress(v) } func (w *parameters) SetPrintRealtime(v bool) { w.p.SetPrintRealtime(v) } func (w *parameters) SetPrintTimestamps(v bool) { w.p.SetPrintTimestamps(v) } +func (w *parameters) SetDebugMode(v bool) { w.p.SetDebugMode(v) } // Diarization (tinydiarize) func (w *parameters) SetDiarize(v bool) { w.p.SetDiarize(v) } diff --git a/bindings/go/pkg/whisper/util_test.go b/bindings/go/pkg/whisper/util_test.go index 8ea2d5b47..1b27255ae 100644 --- a/bindings/go/pkg/whisper/util_test.go +++ b/bindings/go/pkg/whisper/util_test.go @@ -1,6 +1,8 @@ package whisper_test const ( - ModelPath = "../../models/ggml-small.en.bin" - SamplePath = "../../samples/jfk.wav" + ModelPath = "../../models/ggml-small.en.bin" + ModelTinydiarizePath = "../../models/ggml-small.en-tdrz.bin" + SamplePath = "../../samples/jfk.wav" + MultiSpeakerSamplePath = "../../samples/a13.wav" )