Updated go-whisper

Author: djthorpe

pkg/whisper/api/transcribe.go

@@ -13,6 +13,7 @@ import (
 	"github.com/mutablelogic/go-whisper/pkg/whisper"
 	"github.com/mutablelogic/go-whisper/pkg/whisper/segmenter"
 	"github.com/mutablelogic/go-whisper/pkg/whisper/task"
+	"github.com/mutablelogic/go-whisper/pkg/whisper/transcription"
 
 	// Namespace imports
 	. "github.com/djthorpe/go-errors"
@@ -25,11 +26,17 @@ type reqTranscribe struct {
 	File        *multipart.FileHeader `json:"file"`
 	Model       string                `json:"model"`
 	Language    *string               `json:"language"`
-	Prompt      *string               `json:"prompt"`
-	ResponseFmt *string               `json:"response_format"`
 	Temperature *float32              `json:"temperature"`
+	SegmentSize *time.Duration        `json:"segment_size"`
+	ResponseFmt *string               `json:"response_format"`
 }
 
+const (
+	minSegmentSize     = 5 * time.Second
+	maxSegmentSize     = 10 * time.Minute
+	defaultSegmentSize = 5 * time.Minute
+)
+
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
@@ -61,14 +68,15 @@ func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.Respon
 	}
 	defer f.Close()
 
-	// Create a segmenter - read segments of 5 min samples
-	segmenter, err := segmenter.New(f, 5*time.Minute, whisper.SampleRate)
+	// Create a segmenter - read segments based on requested segment size
+	segmenter, err := segmenter.New(f, req.SegmentDur(), whisper.SampleRate)
 	if err != nil {
 		httpresponse.Error(w, http.StatusBadRequest, err.Error())
 		return
 	}
 
 	// Get context for the model, perform transcription
+	var result *transcription.Transcription
 	if err := service.WithModel(model, func(task *task.Context) error {
 		// Check model
 		if translate && !task.CanTranslate() {
@@ -90,18 +98,29 @@ func TranscribeFile(ctx context.Context, service *whisper.Whisper, w http.Respon
 		// 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))
-
+		if err := segmenter.Decode(ctx, func(ts time.Duration, buf []float32) error {
 			// Perform the transcription, return any errors
-			return task.Transcribe(ctx, buf)
-		})
+			return task.Transcribe(ctx, ts, buf, func(segment *transcription.Segment) {
+				fmt.Println("TODO: ", segment)
+			})
+		}); err != nil {
+			return err
+		}
+
+		// End of transcription, get result
+		result = task.Result()
+
+		// Return success
+		return nil
 	}); err != nil {
 		httpresponse.Error(w, http.StatusInternalServerError, err.Error())
 		return
 	}
 
-	var result whisper.Transcription
+	// Set duration
+	result.Duration = segmenter.Duration()
+
+	// Return transcription
 	httpresponse.JSON(w, result, http.StatusOK, 2)
 }
 
@@ -131,3 +150,16 @@ func (r reqTranscribe) ResponseFormat() string {
 	}
 	return *r.ResponseFmt
 }
+
+func (r reqTranscribe) SegmentDur() time.Duration {
+	if r.SegmentSize == nil {
+		return defaultSegmentSize
+	}
+	if *r.SegmentSize < minSegmentSize {
+		return minSegmentSize
+	}
+	if *r.SegmentSize > maxSegmentSize {
+		return maxSegmentSize
+	}
+	return *r.SegmentSize
+}

pkg/whisper/model/store.go

@@ -309,6 +309,9 @@ func listModels(path, ext string) ([]*Model, error) {
 }
 
 func modelNameToId(name string) string {
+	// Lowercase the name, remove the extension
+	name = strings.TrimSuffix(strings.ToLower(name), filepath.Ext(name))
+
 	// We replace all non-alphanumeric characters with underscores
 	return strings.Map(func(r rune) rune {
 		if r >= 'a' && r <= 'z' {

pkg/whisper/segmenter/segmenter.go

@@ -128,3 +128,12 @@ func (s *Segmenter) Decode(ctx context.Context, fn SegmentFunc) error {
 	// Return success
 	return nil
 }
+
+// Return the duration from the file or timestamp
+func (s *Segmenter) Duration() time.Duration {
+	if s.reader != nil {
+		return s.reader.Duration()
+	} else {
+		return s.ts + time.Duration(len(s.buf))*time.Second/time.Duration(s.sample_rate)
+	}
+}

pkg/whisper/task/context.go

@@ -6,9 +6,11 @@ import (
 	"fmt"
 	"path/filepath"
 	"sync"
+	"time"
 
 	// Packages
 	model "github.com/mutablelogic/go-whisper/pkg/whisper/model"
+	transcription "github.com/mutablelogic/go-whisper/pkg/whisper/transcription"
 	whisper "github.com/mutablelogic/go-whisper/sys/whisper"
 
 	// Namespace imports
@@ -22,13 +24,20 @@ import (
 type Context struct {
 	sync.Mutex
 
+	// Model Id and whisper context
 	model   string
 	whisper *whisper.Context
 
 	// Parameters for the next transcription
 	params whisper.FullParams
+
+	// Collect the transcription
+	result *transcription.Transcription
 }
 
+// Callback for new segments during the transcription process
+type NewSegmentFunc func(*transcription.Segment)
+
 //////////////////////////////////////////////////////////////////////////////
 // LIFECYCLE
 
@@ -84,7 +93,6 @@ func (ctx *Context) Close() error {
 
 	// Release resources
 	if ctx.whisper != nil {
-		fmt.Printf("Release model resources %v\n", ctx)
 		whisper.Whisper_free(ctx.whisper)
 	}
 	ctx.whisper = nil
@@ -97,7 +105,7 @@ func (ctx *Context) Close() error {
 //////////////////////////////////////////////////////////////////////////////
 // STRINGIFY
 
-func (ctx Context) MarshalJSON() ([]byte, error) {
+func (ctx *Context) MarshalJSON() ([]byte, error) {
 	type j struct {
 		Model   string             `json:"model"`
 		Params  whisper.FullParams `json:"params"`
@@ -110,7 +118,7 @@ func (ctx Context) MarshalJSON() ([]byte, error) {
 	})
 }
 
-func (ctx Context) String() string {
+func (ctx *Context) String() string {
 	data, err := json.MarshalIndent(ctx, "", "  ")
 	if err != nil {
 		return err.Error()
@@ -145,34 +153,47 @@ func (task *Context) CanTranslate() bool {
 
 // Transcribe samples. The samples should be 16KHz float32 samples in
 // a single channel.
-// TODO: We need a low-latency streaming version of this function.
-// TODO: We need a callback for segment progress.
-func (task *Context) Transcribe(ctx context.Context, samples []float32) error {
+func (task *Context) Transcribe(ctx context.Context, ts time.Duration, samples []float32, fn NewSegmentFunc) error {
 	// Set the 'abort' function
-	/*task.params.SetAbortCallback(task.whisper, func() bool {
+	task.params.SetAbortCallback(task.whisper, func() bool {
 		select {
 		case <-ctx.Done():
 			return true
 		default:
 			return false
 		}
-	})*/
+	})
+
+	// Set the new segment function
+	if fn != nil {
+		task.params.SetSegmentCallback(task.whisper, func(new_segments int) {
+			num_segments := task.whisper.NumSegments()
+			for i := num_segments - new_segments; i < num_segments; i++ {
+				fn(transcription.NewSegment(ts, task.whisper.Segment(i)))
+			}
+		})
+	}
 
-	// Set the 'progress' function
-	//task.params.SetProgressCallback(task.whisper, func(percent int) {
-	//	fmt.Printf("Progress: %v\n", percent)
-	//})
+	// TODO: Set the initial prompt tokens from any previous transcription call
 
 	// Perform the transcription
 	if err := whisper.Whisper_full(task.whisper, task.params, samples); err != nil {
-		return err
+		if ctx.Err() != nil {
+			return ctx.Err()
+		} else {
+			return err
+		}
 	}
 
-	// Get segments
-	for i := 0; i < task.whisper.NumSegments(); i++ {
-		segment := task.whisper.Segment(i)
-		fmt.Printf("Segment: %v\n", segment.Text)
+	// Remove the callbacks
+	task.params.SetAbortCallback(task.whisper, nil)
+	task.params.SetSegmentCallback(task.whisper, nil)
+
+	// Append the transcription
+	if task.result == nil {
+		task.result = transcription.New()
 	}
+	task.result.Append(task.whisper, ts)
 
 	// Return success
 	return nil
@@ -194,6 +215,12 @@ func (ctx *Context) SetLanguage(v string) error {
 	return nil
 }
 
+// Set translate to true or false
 func (ctx *Context) SetTranslate(v bool) {
 	ctx.params.SetTranslate(v)
 }
+
+// Return the transcription result
+func (ctx *Context) Result() *transcription.Transcription {
+	return ctx.result
+}

pkg/whisper/transcription.go

@@ -0,0 +1,81 @@
+package transcription
+
+import (
+	"encoding/json"
+	"time"
+
+	"github.com/mutablelogic/go-whisper/sys/whisper"
+)
+
+//////////////////////////////////////////////////////////////////////////////
+// TYPES
+
+type Transcription struct {
+	Task     string        `json:"task,omitempty"`
+	Language string        `json:"language,omitempty" writer:",width:8"`
+	Duration time.Duration `json:"duration,omitempty" writer:",width:8,right"`
+	Text     string        `json:"text" writer:",width:60,wrap"`
+	Segments []Segment     `json:"segments,omitempty" writer:",width:40,wrap"`
+}
+
+type Segment struct {
+	Id          int32         `json:"id"`
+	Start       time.Duration `json:"start"`
+	End         time.Duration `json:"end"`
+	Text        string        `json:"text"`
+	SpeakerTurn bool          `json:"speaker_turn,omitempty"`
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// LIFECYCLE
+
+// Create a new transcription from a context
+func New() *Transcription {
+	return new(Transcription)
+}
+
+func NewSegment(ts time.Duration, seg *whisper.Segment) *Segment {
+	// Dumb copy function
+	return &Segment{
+		Id:          seg.Id,
+		Text:        seg.Text,
+		Start:       seg.T0 + ts,
+		End:         seg.T1 + ts,
+		SpeakerTurn: seg.SpeakerTurn,
+	}
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// STRINGIFY
+
+func (t *Transcription) String() string {
+	data, err := json.MarshalIndent(t, "", "  ")
+	if err != nil {
+		return err.Error()
+	}
+	return string(data)
+}
+
+func (s *Segment) String() string {
+	data, err := json.MarshalIndent(s, "", "  ")
+	if err != nil {
+		return err.Error()
+	}
+	return string(data)
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS
+
+// Append a transcription to the current transcription, with the offset timestamp
+func (t *Transcription) Append(ctx *whisper.Context, ts time.Duration) {
+	// Append the segment text
+	for i := 0; i < ctx.NumSegments(); i++ {
+		seg := ctx.Segment(i)
+		if t.Text == "" {
+			t.Text = seg.Text
+		} else {
+			t.Text += " " + seg.Text
+		}
+	}
+}