Merge pull request #32 from mutablelogic/v1

V1

Author: djthorpe

pkg/whisper/context/context.go

@@ -2,6 +2,7 @@ package model
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"io/fs"
 	"net/http"
@@ -11,6 +12,7 @@ import (
 	"sync"
 
 	// Packages
+
 	"github.com/mutablelogic/go-whisper/sys/whisper"
 
 	// Namespace imports
@@ -65,6 +67,36 @@ func NewStore(path, ext, modelUrl string) (*Store, error) {
 	return store, nil
 }
 
+//////////////////////////////////////////////////////////////////////////////
+// STRINGIFY
+
+func (s *Store) MarshalJSON() ([]byte, error) {
+	modelNames := func() []string {
+		result := make([]string, len(s.models))
+		for i, model := range s.models {
+			result[i] = model.Id
+		}
+		return result
+	}
+	return json.Marshal(struct {
+		Path   string   `json:"path"`
+		Ext    string   `json:"ext,omitempty"`
+		Models []string `json:"models"`
+	}{
+		Path:   s.path,
+		Ext:    s.ext,
+		Models: modelNames(),
+	})
+}
+
+func (s *Store) String() string {
+	data, err := json.MarshalIndent(s, "", "  ")
+	if err != nil {
+		return err.Error()
+	}
+	return string(data)
+}
+
 //////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 

pkg/whisper/model/store.go

@@ -10,13 +10,18 @@ import (
 
 type opts struct {
 	MaxConcurrent int
+	logfn         LogFn
+	debug         bool
+	gpu           int
 }
 
 type Opt func(*opts) error
+type LogFn func(string)
 
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
+// Set maximum number of concurrent tasks
 func OptMaxConcurrent(v int) Opt {
 	return func(o *opts) error {
 		if v < 1 {
@@ -26,3 +31,27 @@ func OptMaxConcurrent(v int) Opt {
 		return nil
 	}
 }
+
+// Set logging function
+func OptLog(fn LogFn) Opt {
+	return func(o *opts) error {
+		o.logfn = fn
+		return nil
+	}
+}
+
+// Set debugging
+func OptDebug() Opt {
+	return func(o *opts) error {
+		o.debug = true
+		return nil
+	}
+}
+
+// Disable GPU acceleration
+func OptNoGPU() Opt {
+	return func(o *opts) error {
+		o.gpu = -1
+		return nil
+	}
+}

pkg/whisper/opt.go

@@ -1,13 +1,12 @@
 package pool
 
 import (
-
-	// Packages
+	"encoding/json"
 	"fmt"
-	"path/filepath"
 
+	// Packages
 	model "github.com/mutablelogic/go-whisper/pkg/whisper/model"
-	"github.com/mutablelogic/go-whisper/sys/whisper"
+	task "github.com/mutablelogic/go-whisper/pkg/whisper/task"
 
 	// Namespace imports
 	. "github.com/djthorpe/go-errors"
@@ -23,26 +22,25 @@ type ContextPool struct {
 
 	// Base path for models
 	path string
-}
 
-// Context is used for running the transcription or translation
-type Context struct {
-	Model   *model.Model
-	Context *whisper.Context
-	Params  whisper.FullParams
+	// GPU flags
+	gpu int
 }
 
 //////////////////////////////////////////////////////////////////////////////
 // LIFECYCLE
 
 // Create a new context pool of context objects, up to 'max' items
 // Set the path for the model storage
-func NewContextPool(path string, max int32) *ContextPool {
+// If GPU is -1 then disable, if 0 then use default, if >0 then enable
+// and use the specified device
+func NewContextPool(path string, max int, gpu int) *ContextPool {
 	pool := new(ContextPool)
 	pool.Pool = NewPool(max, func() any {
-		return &Context{}
+		return task.New()
 	})
 	pool.path = path
+	pool.gpu = gpu
 
 	// Return success
 	return pool
@@ -53,84 +51,63 @@ func (m *ContextPool) Close() error {
 	return m.Pool.Close()
 }
 
-// Init the context
-func (m *Context) Init(path string, model *model.Model) error {
-	// Check parameters
-	if model == nil {
-		return ErrBadParameter
-	}
-
-	// Get a context
-	ctx := whisper.Whisper_init_from_file_with_params(filepath.Join(path, model.Path), whisper.DefaultContextParams())
-	if ctx == nil {
-		return ErrInternalAppError.With("whisper_init_from_file_with_params")
-	}
-
-	// Set resources
-	m.Context = ctx
-	m.Model = model
-
-	// Return success
-	return nil
+//////////////////////////////////////////////////////////////////////////////
+// STRINGIFY
+
+func (m *ContextPool) MarshalJSON() ([]byte, error) {
+	return json.Marshal(struct {
+		Gpu int `json:"gpu"`
+		N   int `json:"n"`
+		Max int `json:"max"`
+	}{
+		Gpu: m.gpu,
+		N:   m.N(),
+		Max: m.max,
+	})
 }
 
-// Close the context and release all resources
-func (m *Context) Close() error {
-	var result error
-
-	// Do nothing if nil
-	if m == nil {
-		return nil
+func (m *ContextPool) String() string {
+	data, err := json.MarshalIndent(m, "", "  ")
+	if err != nil {
+		return err.Error()
 	}
-
-	// Release resources
-	if m.Context != nil {
-		whisper.Whisper_free(m.Context)
-	}
-	m.Context = nil
-	m.Model = nil
-
-	// Return any errors
-	return result
+	return string(data)
 }
 
 //////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
 // Get a context from the pool, for a model
-func (m *ContextPool) Get(model *model.Model) (*Context, error) {
+func (m *ContextPool) Get(model *model.Model) (*task.Context, error) {
 	// Check parameters
 	if model == nil {
 		return nil, ErrBadParameter
 	}
 
 	// Get a context from the pool
-	ctx, ok := m.Pool.Get().(*Context)
-	if !ok || ctx == nil {
+	t, ok := m.Pool.Get().(*task.Context)
+	if !ok || t == nil {
 		return nil, ErrChannelBlocked.With("unable to get a context from the pool, try again later")
 	}
 
-	// If the model matches, return it
-	if ctx.Model != nil && ctx.Model.Id == model.Id {
-		return ctx, nil
-	}
-
-	// Model didn't match: close the context
-	if err := ctx.Close(); err != nil {
+	// If the model matches, return it, or else release the resources
+	if t.Is(model) {
+		return t, nil
+	} else if err := t.Close(); err != nil {
 		return nil, err
 	}
 
 	// Initialise the context
-	if err := ctx.Init(m.path, model); err != nil {
+	if err := t.Init(m.path, model, m.gpu); err != nil {
 		return nil, err
 	}
 
 	// Return the context
-	return ctx, nil
+	return t, nil
 }
 
 // Put a context back into the pool
-func (m *ContextPool) Put(ctx *Context) {
+func (m *ContextPool) Put(ctx *task.Context) {
 	m.Pool.Put(ctx)
 }
 

pkg/whisper/pool/contextpool.go

@@ -3,12 +3,13 @@ package pool_test
 import (
 	"testing"
 
+	// Packages
 	model "github.com/mutablelogic/go-whisper/pkg/whisper/model"
-	"github.com/mutablelogic/go-whisper/pkg/whisper/pool"
+	pool "github.com/mutablelogic/go-whisper/pkg/whisper/pool"
 )
 
 func Test_contextpool_001(t *testing.T) {
-	var pool = pool.NewContextPool(t.TempDir(), 2)
+	var pool = pool.NewContextPool(t.TempDir(), 2, 0)
 
 	model1, err := pool.Get(&model.Model{
 		Id: "model1",