refactor(go bindings): add diarization unit tests

bindings/go/.gitignore

@@ -1,2 +1,3 @@
 build
 models
+samples/a13.wav

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}

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 {

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)

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

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)
 

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)
 

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) }

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"
 )