Integrated segmenter

Author: djthorpe

pkg/whisper/api/transcribe.go

@@ -5,13 +5,13 @@ import (
 	"fmt"
 	"mime/multipart"
 	"net/http"
+	"time"
 
 	// Packages
-
-	"github.com/go-audio/wav"
 	"github.com/mutablelogic/go-server/pkg/httprequest"
 	"github.com/mutablelogic/go-server/pkg/httpresponse"
 	"github.com/mutablelogic/go-whisper/pkg/whisper"
+	"github.com/mutablelogic/go-whisper/pkg/whisper/segmenter"
 	"github.com/mutablelogic/go-whisper/pkg/whisper/task"
 
 	// Namespace imports
@@ -53,12 +53,6 @@ func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.Respon
 		return
 	}
 
-	// Check audio format - allow WAV or binary
-	if req.File.Header.Get("Content-Type") != "audio/wav" && req.File.Header.Get("Content-Type") != httprequest.ContentTypeBinary {
-		httpresponse.Error(w, http.StatusBadRequest, "unsupported audio format:", req.File.Header.Get("Content-Type"))
-		return
-	}
-
 	// Open file
 	f, err := req.File.Open()
 	if err != nil {
@@ -67,15 +61,14 @@ func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.Respon
 	}
 	defer f.Close()
 
-	// Read samples
-	buf, err := wav.NewDecoder(f).FullPCMBuffer()
+	// Create a segmenter - read segments of 5 min samples
+	segmenter, err := segmenter.New(f, 5*time.Minute, whisper.SampleRate)
 	if err != nil {
-		httpresponse.Error(w, http.StatusInternalServerError, err.Error())
+		httpresponse.Error(w, http.StatusBadRequest, err.Error())
 		return
 	}
 
 	// Get context for the model, perform transcription
-	var result *whisper.Transcription
 	if err := service.WithModel(model, func(task *task.Context) error {
 		// Check model
 		if translate && !task.CanTranslate() {
@@ -93,29 +86,23 @@ func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.Respon
 				return err
 			}
 		}
-		// TODO Set prompt and temperature
-		/*
-			if req.Prompt != nil {
-				ctx.SetPrompt(*req.Prompt)
-			}
-			if req.Temperature != nil {
-				ctx.SetTemperature(*req.Temperature)
-			}
-		*/
-		// Perform the transcription, return any errors
-		return task.Transcribe(ctx, buf.AsFloat32Buffer().Data)
+
+		// TODO: Set temperature, etc
+
+		// Read samples and transcribe them
+		return segmenter.Decode(ctx, func(ts time.Duration, buf []float32) error {
+			fmt.Println("audio segment", ts, len(buf))
+
+			// Perform the transcription, return any errors
+			return task.Transcribe(ctx, buf)
+		})
 	}); err != nil {
-		httpresponse.Error(w, http.StatusBadRequest, err.Error())
+		httpresponse.Error(w, http.StatusInternalServerError, err.Error())
 		return
 	}
 
-	// Response - TODO srt, vtt, verbose_json
-	switch req.ResponseFormat() {
-	case "text":
-		httpresponse.Text(w, result.Text, http.StatusOK)
-	default:
-		httpresponse.JSON(w, result, http.StatusOK, 2)
-	}
+	var result whisper.Transcription
+	httpresponse.JSON(w, result, http.StatusOK, 2)
 }
 
 ///////////////////////////////////////////////////////////////////////////////

pkg/whisper/segmenter/segmenter.go

@@ -24,14 +24,14 @@ type Segmenter struct {
 
 // SegmentFunc is a callback function which is called when a segment is ready
 // to be processed. The first argument is the timestamp of the segment.
-type SegmentFunc func(time.Duration, []float32)
+type SegmentFunc func(time.Duration, []float32) error
 
 //////////////////////////////////////////////////////////////////////////////
 // LIFECYCLE
 
 // Create a new segmenter for "NumSamples" with a reader r
 // If NumSamples is zero then no segmenting is performed
-func NewSegmenter(r io.Reader, dur time.Duration, sample_rate int) (*Segmenter, error) {
+func New(r io.Reader, dur time.Duration, sample_rate int) (*Segmenter, error) {
 	segmenter := new(Segmenter)
 
 	// Check arguments
@@ -105,7 +105,9 @@ func (s *Segmenter) Decode(ctx context.Context, fn SegmentFunc) error {
 
 		// n != 0 and len(buf) >= n we have a segment to process
 		if s.n != 0 && len(s.buf) >= s.n {
-			fn(s.ts, s.buf)
+			if err := fn(s.ts, s.buf); err != nil {
+				return err
+			}
 			// Increment the timestamp
 			s.ts += time.Duration(len(s.buf)) * time.Second / time.Duration(s.sample_rate)
 			// Clear the buffer

pkg/whisper/segmenter/segmenter_test.go

@@ -20,14 +20,15 @@ func Test_segmenter_001(t *testing.T) {
 	if !assert.NoError(err) {
 		t.SkipNow()
 	}
-	segmenter, err := segmenter.NewSegmenter(f, 200*time.Millisecond, 16000)
+	segmenter, err := segmenter.New(f, 200*time.Millisecond, 16000)
 	if !assert.NoError(err) {
 		t.SkipNow()
 	}
 	defer segmenter.Close()
 
-	assert.NoError(segmenter.Decode(context.Background(), func(ts time.Duration, buf []float32) {
+	assert.NoError(segmenter.Decode(context.Background(), func(ts time.Duration, buf []float32) error {
 		t.Log(ts, len(buf))
+		return nil
 	}))
 }
 
@@ -40,13 +41,14 @@ func Test_segmenter_002(t *testing.T) {
 	}
 
 	// No segmentation, just output the audio
-	segmenter, err := segmenter.NewSegmenter(f, 0, 16000)
+	segmenter, err := segmenter.New(f, 0, 16000)
 	if !assert.NoError(err) {
 		t.SkipNow()
 	}
 	defer segmenter.Close()
 
-	assert.NoError(segmenter.Decode(context.Background(), func(ts time.Duration, buf []float32) {
+	assert.NoError(segmenter.Decode(context.Background(), func(ts time.Duration, buf []float32) error {
 		t.Log(ts, len(buf))
+		return nil
 	}))
 }

pkg/whisper/whisper.go

@@ -36,6 +36,9 @@ const (
 
 	// This is where the model is downloaded from
 	defaultModelUrl = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/?download=true"
+
+	// Sample Rate
+	SampleRate = whisper.SampleRate
 )
 
 //////////////////////////////////////////////////////////////////////////////

sys/whisper/whisper.go

@@ -19,6 +19,13 @@ type (
 	Context C.struct_whisper_context
 )
 
+///////////////////////////////////////////////////////////////////////////////
+// Constants
+
+const (
+	SampleRate = C.WHISPER_SAMPLE_RATE
+)
+
 ///////////////////////////////////////////////////////////////////////////////
 // LIFECYCLE