feat(whisper): Add diarization (tinydiarize)

Signed-off-by: Richard Palethorpe <[email protected]>

Author: richiejp

backend/backend.proto

@@ -276,6 +276,7 @@ message TranscriptRequest {
   string language = 3;
   uint32 threads = 4;
   bool translate = 5;
+  bool diarize = 6;
 }
 
 message TranscriptResult {
@@ -305,7 +306,7 @@ message GenerateImageRequest {
   // Diffusers
   string EnableParameters = 10;
   int32 CLIPSkip = 11;
-  
+
   // Reference images for models that support them (e.g., Flux Kontext)
   repeated string ref_images = 12;
 }

backend/go/whisper/gowhisper.cpp

@@ -7,34 +7,35 @@ static struct whisper_vad_context *vctx;
 static struct whisper_context *ctx;
 static std::vector<float> flat_segs;
 
-static void ggml_log_cb(enum ggml_log_level level, const char* log, void* data) {
-    const char* level_str;
-
-    if (!log) {
-        return;
-    }
-
-    switch (level) {
-        case GGML_LOG_LEVEL_DEBUG:
-            level_str = "DEBUG";
-            break;
-        case GGML_LOG_LEVEL_INFO:
-            level_str = "INFO";
-            break;
-        case GGML_LOG_LEVEL_WARN:
-            level_str = "WARN";
-            break;
-        case GGML_LOG_LEVEL_ERROR:
-            level_str = "ERROR";
-            break;
-        default: /* Potential future-proofing */
-            level_str = "?????";
-            break;
-    }
-
-    fprintf(stderr, "[%-5s] ", level_str);
-    fputs(log, stderr);
-    fflush(stderr);
+static void ggml_log_cb(enum ggml_log_level level, const char *log,
+                        void *data) {
+  const char *level_str;
+
+  if (!log) {
+    return;
+  }
+
+  switch (level) {
+  case GGML_LOG_LEVEL_DEBUG:
+    level_str = "DEBUG";
+    break;
+  case GGML_LOG_LEVEL_INFO:
+    level_str = "INFO";
+    break;
+  case GGML_LOG_LEVEL_WARN:
+    level_str = "WARN";
+    break;
+  case GGML_LOG_LEVEL_ERROR:
+    level_str = "ERROR";
+    break;
+  default: /* Potential future-proofing */
+    level_str = "?????";
+    break;
+  }
+
+  fprintf(stderr, "[%-5s] ", level_str);
+  fputs(log, stderr);
+  fflush(stderr);
 }
 
 int load_model(const char *const model_path) {
@@ -105,8 +106,8 @@ int vad(float pcmf32[], size_t pcmf32_len, float **segs_out,
   return 0;
 }
 
-int transcribe(uint32_t threads, char *lang, bool translate, float pcmf32[],
-               size_t pcmf32_len, size_t *segs_out_len) {
+int transcribe(uint32_t threads, char *lang, bool translate, bool tdrz,
+               float pcmf32[], size_t pcmf32_len, size_t *segs_out_len) {
   whisper_full_params wparams =
       whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
 
@@ -120,6 +121,9 @@ int transcribe(uint32_t threads, char *lang, bool translate, float pcmf32[],
   wparams.translate = translate;
   wparams.debug_mode = true;
   wparams.print_progress = true;
+  wparams.tdrz_enable = tdrz;
+
+  fprintf(stderr, "info: Enable tdrz: %d\n", tdrz);
 
   if (whisper_full(ctx, wparams, pcmf32, pcmf32_len)) {
     fprintf(stderr, "error: transcription failed\n");
@@ -144,3 +148,7 @@ int n_tokens(int i) { return whisper_full_n_tokens(ctx, i); }
 int32_t get_token_id(int i, int j) {
   return whisper_full_get_token_id(ctx, i, j);
 }
+
+bool get_segment_speaker_turn_next(int i) {
+  return whisper_full_get_segment_speaker_turn_next(ctx, i);
+}

backend/go/whisper/gowhisper.go

@@ -14,15 +14,16 @@ import (
 )
 
 var (
-	CppLoadModel       func(modelPath string) int
-	CppLoadModelVAD    func(modelPath string) int
-	CppVAD             func(pcmf32 []float32, pcmf32Size uintptr, segsOut unsafe.Pointer, segsOutLen unsafe.Pointer) int
-	CppTranscribe      func(threads uint32, lang string, translate bool, pcmf32 []float32, pcmf32Len uintptr, segsOutLen unsafe.Pointer) int
-	CppGetSegmentText  func(i int) string
-	CppGetSegmentStart func(i int) int64
-	CppGetSegmentEnd   func(i int) int64
-	CppNTokens         func(i int) int
-	CppGetTokenID      func(i int, j int) int
+	CppLoadModel                 func(modelPath string) int
+	CppLoadModelVAD              func(modelPath string) int
+	CppVAD                       func(pcmf32 []float32, pcmf32Size uintptr, segsOut unsafe.Pointer, segsOutLen unsafe.Pointer) int
+	CppTranscribe                func(threads uint32, lang string, translate bool, diarize bool, pcmf32 []float32, pcmf32Len uintptr, segsOutLen unsafe.Pointer) int
+	CppGetSegmentText            func(i int) string
+	CppGetSegmentStart           func(i int) int64
+	CppGetSegmentEnd             func(i int) int64
+	CppNTokens                   func(i int) int
+	CppGetTokenID                func(i int, j int) int
+	CppGetSegmentSpeakerTurnNext func(i int) bool
 )
 
 type Whisper struct {
@@ -122,7 +123,7 @@ func (w *Whisper) AudioTranscription(opts *pb.TranscriptRequest) (pb.TranscriptR
 	segsLen := uintptr(0xdeadbeef)
 	segsLenPtr := unsafe.Pointer(&segsLen)
 
-	if ret := CppTranscribe(opts.Threads, opts.Language, opts.Translate, data, uintptr(len(data)), segsLenPtr); ret != 0 {
+	if ret := CppTranscribe(opts.Threads, opts.Language, opts.Translate, opts.Diarize, data, uintptr(len(data)), segsLenPtr); ret != 0 {
 		return pb.TranscriptResult{}, fmt.Errorf("Failed Transcribe")
 	}
 
@@ -134,6 +135,10 @@ func (w *Whisper) AudioTranscription(opts *pb.TranscriptRequest) (pb.TranscriptR
 		txt := strings.Clone(CppGetSegmentText(i))
 		tokens := make([]int32, CppNTokens(i))
 
+		if opts.Diarize && CppGetSegmentSpeakerTurnNext(i) {
+			txt += " [SPEAKER_TURN]"
+		}
+
 		for j := range tokens {
 			tokens[j] = int32(CppGetTokenID(i, j))
 		}
@@ -151,6 +156,6 @@ func (w *Whisper) AudioTranscription(opts *pb.TranscriptRequest) (pb.TranscriptR
 
 	return pb.TranscriptResult{
 		Segments: segments,
-		Text: strings.TrimSpace(text),
+		Text:     strings.TrimSpace(text),
 	}, nil
 }

backend/go/whisper/gowhisper.h

@@ -6,11 +6,12 @@ int load_model(const char *const model_path);
 int load_model_vad(const char *const model_path);
 int vad(float pcmf32[], size_t pcmf32_size, float **segs_out,
         size_t *segs_out_len);
-int transcribe(uint32_t threads, char *lang, bool translate, float pcmf32[],
-               size_t pcmf32_len, size_t *segs_out_len);
+int transcribe(uint32_t threads, char *lang, bool translate, bool tdrz,
+               float pcmf32[], size_t pcmf32_len, size_t *segs_out_len);
 const char *get_segment_text(int i);
 int64_t get_segment_t0(int i);
 int64_t get_segment_t1(int i);
 int n_tokens(int i);
 int32_t get_token_id(int i, int j);
+bool get_segment_speaker_turn_next(int i);
 }

backend/go/whisper/main.go

@@ -33,6 +33,7 @@ func main() {
 		{&CppGetSegmentEnd, "get_segment_t1"},
 		{&CppNTokens, "n_tokens"},
 		{&CppGetTokenID, "get_token_id"},
+		{&CppGetSegmentSpeakerTurnNext, "get_segment_speaker_turn_next"},
 	}
 
 	for _, lf := range libFuncs {

core/backend/transcript.go

@@ -12,7 +12,7 @@ import (
 	"github.com/mudler/LocalAI/pkg/model"
 )
 
-func ModelTranscription(audio, language string, translate bool, ml *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (*schema.TranscriptionResult, error) {
+func ModelTranscription(audio, language string, translate bool, diarize bool, ml *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (*schema.TranscriptionResult, error) {
 
 	if modelConfig.Backend == "" {
 		modelConfig.Backend = model.WhisperBackend
@@ -34,6 +34,7 @@ func ModelTranscription(audio, language string, translate bool, ml *model.ModelL
 		Dst:       audio,
 		Language:  language,
 		Translate: translate,
+		Diarize:   diarize,
 		Threads:   uint32(*modelConfig.Threads),
 	})
 	if err != nil {

core/cli/transcript.go

@@ -20,6 +20,7 @@ type TranscriptCMD struct {
 	Model      string `short:"m" required:"" help:"Model name to run the TTS"`
 	Language   string `short:"l" help:"Language of the audio file"`
 	Translate  bool   `short:"c" help:"Translate the transcription to english"`
+	Diarize    bool   `short:"d" help:"Mark speaker turns"`
 	Threads    int    `short:"t" default:"1" help:"Number of threads used for parallel computation"`
 	ModelsPath string `env:"LOCALAI_MODELS_PATH,MODELS_PATH" type:"path" default:"${basepath}/models" help:"Path containing models used for inferencing" group:"storage"`
 }
@@ -56,7 +57,7 @@ func (t *TranscriptCMD) Run(ctx *cliContext.Context) error {
 		}
 	}()
 
-	tr, err := backend.ModelTranscription(t.Filename, t.Language, t.Translate, ml, c, opts)
+	tr, err := backend.ModelTranscription(t.Filename, t.Language, t.Translate, t.Diarize, ml, c, opts)
 	if err != nil {
 		return err
 	}

core/http/endpoints/openai/transcription.go

@@ -36,6 +36,8 @@ func TranscriptEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, app
 			return fiber.ErrBadRequest
 		}
 
+		diarize := c.FormValue("diarize", "false") != "false"
+
 		// retrieve the file data from the request
 		file, err := c.FormFile("file")
 		if err != nil {
@@ -67,7 +69,7 @@ func TranscriptEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, app
 
 		log.Debug().Msgf("Audio file copied to: %+v", dst)
 
-		tr, err := backend.ModelTranscription(dst, input.Language, input.Translate, ml, *config, appConfig)
+		tr, err := backend.ModelTranscription(dst, input.Language, input.Translate, diarize, ml, *config, appConfig)
 		if err != nil {
 			return err
 		}

gallery/index.yaml

@@ -20174,6 +20174,15 @@
     - filename: "ggml-small.bin"
       uri: "huggingface://ggerganov/whisper.cpp/ggml-small.bin"
       sha256: 1be3a9b2063867b937e64e2ec7483364a79917e157fa98c5d94b5c1fffea987b
+- !!merge <<: *whisper
+  name: "whisper-small-en-tdrz"
+  overrides:
+    parameters:
+      model: ggml-small.en-tdrz.bin
+  files:
+    - filename: "ggml-small.bin"
+      uri: "huggingface://akashmjn/tinydiarize-whisper.cpp/ggml-small.en-tdrz.bin"
+      sha256: ceac3ec06d1d98ef71aec665283564631055fd6129b79d8e1be4f9cc33cc54b4
 - !!merge <<: *whisper
   name: "whisper-small-en-q5_1"
   overrides: