Updated

Author: djthorpe

pkg/whisper/task/context.go

@@ -1,6 +1,7 @@
 package task
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"path/filepath"
@@ -20,7 +21,9 @@ import (
 type Context struct {
 	model   string
 	whisper *whisper.Context
-	params  whisper.FullParams
+
+	// Parameters for the next transcription
+	params whisper.FullParams
 }
 
 //////////////////////////////////////////////////////////////////////////////
@@ -60,8 +63,6 @@ func (m *Context) Init(path string, model *model.Model, gpu int) error {
 	// Set resources
 	m.whisper = ctx
 	m.model = model.Id
-	m.params = whisper.DefaultFullParams(whisper.SAMPLING_GREEDY)
-	m.params.SetLanguage("auto")
 
 	// Return success
 	return nil
@@ -125,13 +126,39 @@ func (ctx *Context) Is(model *model.Model) bool {
 	return ctx.model == model.Id
 }
 
+// Copy task parameters from the default
+func (task *Context) CopyParams() {
+	task.params = whisper.DefaultFullParams(whisper.SAMPLING_GREEDY)
+	task.params.SetLanguage("auto")
+}
+
 // 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 (ctx *Context) Transcribe(samples []float32) error {
+func (task *Context) Transcribe(ctx context.Context, samples []float32) error {
+	// Set the 'abort' function
+	task.params.SetAbortCallback(task.whisper, func() bool {
+		select {
+		case <-ctx.Done():
+			return true
+		default:
+			return false
+		}
+	})
+
+	// Set the 'progress' function
+	task.params.SetProgressCallback(task.whisper, func(percent int) {
+		fmt.Printf("Progress: %v\n", percent)
+	})
+
 	// Perform the transcription
-	return whisper.Whisper_full(ctx.whisper, ctx.params, samples)
+	if err := whisper.Whisper_full(task.whisper, task.params, samples); err != nil {
+		return err
+	}
+
+	// Return success
+	return nil
 }
 
 // Set the language. For transcription, this is the language of the

pkg/whisper/whisper.go

@@ -175,6 +175,9 @@ func (w *Whisper) WithModel(model *model.Model, fn func(task *task.Context) erro
 	}
 	defer w.pool.Put(task)
 
+	// Copy parameters
+	task.CopyParams()
+
 	// Execute the function
 	return fn(task)
 }

pkg/whisper/whisper_test.go

@@ -225,7 +225,7 @@ func Test_whisper_005(t *testing.T) {
 
 		assert.NoError(service.WithModel(model, func(task *task.Context) error {
 			t.Log("Transcribing", len(samples), "samples")
-			return task.Transcribe(samples)
+			return task.Transcribe(context.Background(), samples)
 		}))
 	})
 
@@ -241,7 +241,7 @@ func Test_whisper_005(t *testing.T) {
 
 		assert.NoError(service.WithModel(model, func(task *task.Context) error {
 			t.Log("Transcribing", len(samples), "samples")
-			return task.Transcribe(samples)
+			return task.Transcribe(context.Background(), samples)
 		}))
 	})
 
@@ -257,7 +257,7 @@ func Test_whisper_005(t *testing.T) {
 
 		assert.NoError(service.WithModel(model, func(task *task.Context) error {
 			t.Log("Transcribing", len(samples), "samples")
-			return task.Transcribe(samples)
+			return task.Transcribe(context.Background(), samples)
 		}))
 	})
 }
@@ -294,7 +294,7 @@ func Test_whisper_006(t *testing.T) {
 
 			assert.NoError(service.WithModel(model, func(task *task.Context) error {
 				t.Log("Transcribing", len(samples), "samples")
-				return task.Transcribe(samples)
+				return task.Transcribe(context.Background(), samples)
 			}))
 		})
 
@@ -312,7 +312,7 @@ func Test_whisper_006(t *testing.T) {
 
 			assert.NoError(service.WithModel(model, func(task *task.Context) error {
 				t.Log("Transcribing", len(samples), "samples")
-				return task.Transcribe(samples)
+				return task.Transcribe(context.Background(), samples)
 			}))
 		})
 
@@ -330,7 +330,7 @@ func Test_whisper_006(t *testing.T) {
 
 			assert.NoError(service.WithModel(model, func(task *task.Context) error {
 				t.Log("Transcribing", len(samples), "samples")
-				return task.Transcribe(samples)
+				return task.Transcribe(context.Background(), samples)
 			}))
 		})
 	})