Merge pull request #52 from mutablelogic/v1

Add diarization and output formats

Author: djthorpe

Makefile

@@ -56,7 +56,7 @@ docker: docker-dep submodule
 		--build-arg OS=${OS} \
 		--build-arg SOURCE=${BUILD_MODULE} \
 		--build-arg VERSION=${VERSION} \
-		-f etc/Dockerfile.${ARCH} .
+		-f etc/Dockerfile.${OS}-${ARCH} .
 
 # Test whisper bindings
 test: generate libwhisper libggml

README.md

@@ -62,7 +62,7 @@ curl -F model=ggml-medium-q5_0 -F file=@samples/jfk.wav localhost:8080/v1/audio/
 To translate a media file into a different language, you can use the following command:
 
 ```bash
-curl -F model=ggml-medium-q5_0 -F file=@samples/ge-podcast.wav -F language=en localhost:8080/v1/audio/translations\?stream=true
+curl -F model=ggml-medium-q5_0 -F file=@samples/de-podcast.wav -F language=en localhost:8080/v1/audio/translations\?stream=true
 ```
 
 There's more information on the API [here](doc/API.md).

cmd/server/flags.go

@@ -20,9 +20,9 @@ func NewFlags(name string, args []string) (*Flags, error) {
 		FlagSet: flag.NewFlagSet(name, flag.ContinueOnError),
 	}
 	flags.endpoint = flags.String("endpoint", "/v1", "HTTP endpoint")
-	flags.listen = flags.String("listen", ":8080", "HTTP Listen address")
+	flags.listen = flags.String("listen", "127.0.0.1:8080", "HTTP Listen address")
 	flags.dir = flags.String("dir", "${WHISPER_DATA}", "Model data directory")
-	flags.debug = flags.Bool("debug", false, "Display debug information")
+	flags.debug = flags.Bool("debug", false, "Output additional debug information")
 
 	// Parse flags and return any error
 	return flags, flags.Parse(args)

cmd/server/main.go

@@ -15,6 +15,7 @@ import (
 	httpserver "github.com/mutablelogic/go-server/pkg/httpserver"
 	whisper "github.com/mutablelogic/go-whisper/pkg/whisper"
 	api "github.com/mutablelogic/go-whisper/pkg/whisper/api"
+	version "github.com/mutablelogic/go-whisper/pkg/whisper/version"
 )
 
 func main() {
@@ -45,6 +46,16 @@ func main() {
 		os.Exit(-1)
 	}
 
+	// Print version
+	if version.GitSource != "" {
+		log.Println(name, version.GitSource)
+	} else {
+		log.Println(name)
+	}
+	if version.GitTag != "" {
+		log.Println("Version:", version.GitTag)
+	}
+
 	// Create a whisper service
 	log.Println("Storing models at", dir)
 	opts := []whisper.Opt{
@@ -77,7 +88,7 @@ func main() {
 	api.RegisterEndpoints(flags.Endpoint(), mux, whisper)
 
 	// Create a new HTTP server
-	log.Println("List address", flags.Listen())
+	log.Println("Listen address", flags.Listen())
 	server, err := httpserver.Config{
 		Listen: flags.Listen(),
 		Router: mux,

doc/API.md

@@ -144,9 +144,20 @@ event: ok
 
 ### Translation
 
-This is the same as transcription (above) except that the `language` parameter is not optional, and should be the language to translate the audio into.
+This is the same as transcription (above) except that the `language` parameter is always set to 'en', to translate the audio into English.
 
 ```html
 POST /v1/audio/translations
 POST /v1/audio/translations?stream={bool}
 ```
+
+### Diarization
+
+To diarize an Enlgish-language audio file, use the following endpoint:
+
+```html
+POST /v1/audio/diarize
+POST /v1/audio/diarize?stream={bool}
+```
+
+The segments returned include a "speaker_turn" field which indicates that the segment is a new speaker. It requires a separate download of a [diarization model](https://huggingface.co/akashmjn/tinydiarize-whisper.cpp).

etc/Dockerfile.amd64 renamed to etc/Dockerfile.linux-amd64

@@ -70,7 +70,7 @@ func RegisterEndpoints(base string, mux *http.ServeMux, whisper *whisper.Whisper
 
 		switch r.Method {
 		case http.MethodPost:
-			TranscribeFile(r.Context(), whisper, w, r, true)
+			TranscribeFile(r.Context(), whisper, w, r, Translate)
 		default:
 			httpresponse.Error(w, http.StatusMethodNotAllowed)
 		}
@@ -84,7 +84,21 @@ func RegisterEndpoints(base string, mux *http.ServeMux, whisper *whisper.Whisper
 
 		switch r.Method {
 		case http.MethodPost:
-			TranscribeFile(r.Context(), whisper, w, r, false)
+			TranscribeFile(r.Context(), whisper, w, r, Transcribe)
+		default:
+			httpresponse.Error(w, http.StatusMethodNotAllowed)
+		}
+	})
+
+	// Diarize: POST /v1/audio/diarize
+	//   Transcribes audio into the input language - language parameter should be set to the source
+	//   language of the audio. Output speaker parts.
+	mux.HandleFunc(joinPath(base, "audio/diarize"), func(w http.ResponseWriter, r *http.Request) {
+		defer r.Body.Close()
+
+		switch r.Method {
+		case http.MethodPost:
+			TranscribeFile(r.Context(), whisper, w, r, Diarize)
 		default:
 			httpresponse.Error(w, http.StatusMethodNotAllowed)
 		}

etc/Dockerfile.arm64 renamed to etc/Dockerfile.linux-arm64

@@ -1,10 +1,12 @@
 package api
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"mime/multipart"
 	"net/http"
+	"strings"
 	"time"
 
 	// Packages
@@ -35,16 +37,37 @@ type queryTranscribe struct {
 	Stream bool `json:"stream"`
 }
 
+type TaskType int
+type ResponseFormat string
+
+///////////////////////////////////////////////////////////////////////////////
+// GLOBALS
+
 const (
 	minSegmentSize     = 5 * time.Second
 	maxSegmentSize     = 10 * time.Minute
 	defaultSegmentSize = 5 * time.Minute
 )
 
+const (
+	_          TaskType = iota
+	Transcribe          // Transcribe audio
+	Translate           // Translate text
+	Diarize             // Diarize audio
+)
+
+const (
+	FormatJson        ResponseFormat = "json"
+	FormatText        ResponseFormat = "text"
+	FormatSrt         ResponseFormat = "srt"
+	FormatVerboseJson ResponseFormat = "verbose_json"
+	FormatVtt         ResponseFormat = "vtt"
+)
+
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
-func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.ResponseWriter, r *http.Request, translate bool) {
+func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.ResponseWriter, r *http.Request, t TaskType) {
 	var req reqTranscribe
 	var query queryTranscribe
 	if err := httprequest.Query(&query, r.URL.Query()); err != nil {
@@ -96,54 +119,86 @@ func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.Respon
 
 	// Get context for the model, perform transcription
 	var result *schema.Transcription
-	if err := service.WithModel(model, func(task *task.Context) error {
-		// Check model
-		if translate && !task.CanTranslate() {
-			return ErrBadParameter.With("model is not multilingual, cannot translate")
-		}
+	if err := service.WithModel(model, func(taskctx *task.Context) error {
+		result = taskctx.Result()
 
-		// Set parameters for transcription & translation, default to english
-		task.SetTranslate(translate)
-		if req.Language != nil {
-			if err := task.SetLanguage(*req.Language); err != nil {
-				return err
+		switch t {
+		case Translate:
+			// Check model
+			if !taskctx.CanTranslate() {
+				return ErrBadParameter.With("model is not multilingual, cannot translate")
 			}
-		} else if translate {
-			if err := task.SetLanguage("en"); err != nil {
+			taskctx.SetTranslate(true)
+			taskctx.SetDiarize(false)
+			result.Task = "translate"
+
+			// Set language to EN
+			if err := taskctx.SetLanguage("en"); err != nil {
 				return err
 			}
+		case Diarize:
+			taskctx.SetTranslate(false)
+			taskctx.SetDiarize(true)
+			result.Task = "diarize"
+
+			// Set language
+			if req.Language != nil {
+				if err := taskctx.SetLanguage(*req.Language); err != nil {
+					return err
+				}
+			}
+		default:
+			// Transcribe
+			taskctx.SetTranslate(false)
+			taskctx.SetDiarize(false)
+			result.Task = "transribe"
+
+			// Set language
+			if req.Language != nil {
+				if err := taskctx.SetLanguage(*req.Language); err != nil {
+					return err
+				}
+			}
 		}
 
 		// TODO: Set temperature, etc
 
 		// Create response
-		result = task.Result()
-		result.Task = "transcribe"
-		if translate {
-			result.Task = "translate"
-		}
-		result.Duration = schema.Timestamp(segmenter.Duration())
-		result.Language = task.Language()
+		result.Language = taskctx.Language()
 
 		// Output the header
 		if stream != nil {
-			stream.Write("task", result)
+			stream.Write("task", taskctx.Result())
 		}
 
 		// Read samples and transcribe them
 		if err := segmenter.Decode(ctx, func(ts time.Duration, buf []float32) error {
 			// Perform the transcription, return any errors
-			return task.Transcribe(ctx, ts, buf, req.OutputSegments() || stream != nil, func(segment *schema.Segment) {
-				if stream != nil {
+			return taskctx.Transcribe(ctx, ts, buf, func(segment *schema.Segment) {
+				// Segment callback
+				if stream == nil {
+					return
+				}
+				var buf bytes.Buffer
+				switch req.ResponseFormat() {
+				case FormatVerboseJson, FormatJson:
 					stream.Write("segment", segment)
+					return
+				case FormatSrt:
+					task.WriteSegmentSrt(&buf, segment)
+				case FormatVtt:
+					task.WriteSegmentVtt(&buf, segment)
+				case FormatText:
+					task.WriteSegmentText(&buf, segment)
 				}
+				stream.Write("segment", buf.String())
 			})
 		}); err != nil {
 			return err
 		}
 
 		// Set the language
-		result.Language = task.Language()
+		result.Language = taskctx.Language()
 
 		// Return success
 		return nil
@@ -156,11 +211,32 @@ func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.Respon
 		return
 	}
 
-	// Return transcription if not streaming
-	if stream == nil {
-		httpresponse.JSON(w, result, http.StatusOK, 2)
-	} else {
+	// Return streaming ok
+	if stream != nil {
 		stream.Write("ok")
+		return
+	}
+
+	// Rrturn result based on response format
+	switch req.ResponseFormat() {
+	case FormatJson, FormatVerboseJson:
+		httpresponse.JSON(w, result, http.StatusOK, 0)
+	case FormatText:
+		httpresponse.Text(w, "", http.StatusOK)
+		for _, seg := range result.Segments {
+			task.WriteSegmentText(w, seg)
+		}
+		w.Write([]byte("\n"))
+	case FormatSrt:
+		httpresponse.Text(w, "", http.StatusOK, "Content-Type", "application/x-subrip")
+		for _, seg := range result.Segments {
+			task.WriteSegmentSrt(w, seg)
+		}
+	case FormatVtt:
+		httpresponse.Text(w, "WEBVTT\n\n", http.StatusOK, "Content-Type", "text/vtt")
+		for _, seg := range result.Segments {
+			task.WriteSegmentVtt(w, seg)
+		}
 	}
 }
 
@@ -184,17 +260,29 @@ func (r reqTranscribe) Validate() error {
 	}
 	return nil
 }
-func (r reqTranscribe) ResponseFormat() string {
+func (r reqTranscribe) ResponseFormat() ResponseFormat {
 	if r.ResponseFmt == nil {
-		return "json"
+		return FormatJson
+	}
+	switch strings.ToLower(*r.ResponseFmt) {
+	case "json":
+		return FormatJson
+	case "text":
+		return FormatText
+	case "srt":
+		return FormatSrt
+	case "verbose_json":
+		return FormatVerboseJson
+	case "vtt":
+		return FormatVtt
 	}
-	return *r.ResponseFmt
+	return FormatJson
 }
 
 func (r reqTranscribe) OutputSegments() bool {
 	// We want to output segments if the response format is  "srt", "verbose_json", "vtt"
 	switch r.ResponseFormat() {
-	case "srt", "verbose_json", "vtt":
+	case FormatSrt, FormatVerboseJson, FormatVtt:
 		return true
 	default:
 		return false