From b83b27c3a006cd38e887799d4a82278287fe80ab Mon Sep 17 00:00:00 2001 From: David Thorpe Date: Thu, 1 Aug 2024 13:15:45 +0200 Subject: [PATCH] Reorganized code --- Makefile | 4 +- cmd/server/main.go | 2 +- cmd/whisper/download.go | 22 ++++ cmd/whisper/main.go | 101 ++++++++++++++++++ cmd/whisper/models.go | 12 +++ cmd/whisper/server.go | 37 +++++++ pkg/whisper/opt.go => opt.go | 0 pkg/whisper/api/models.go | 2 +- pkg/whisper/api/register.go | 2 +- pkg/whisper/api/transcribe.go | 2 +- pkg/whisper/pool/contextpool.go | 3 +- pkg/whisper/whisper.go => whisper.go | 4 +- .../whisper_test.go => whisper_test.go | 8 +- 13 files changed, 185 insertions(+), 14 deletions(-) create mode 100644 cmd/whisper/download.go create mode 100644 cmd/whisper/main.go create mode 100644 cmd/whisper/models.go create mode 100644 cmd/whisper/server.go rename pkg/whisper/opt.go => opt.go (100%) rename pkg/whisper/whisper.go => whisper.go (97%) rename pkg/whisper/whisper_test.go => whisper_test.go (97%) diff --git a/Makefile b/Makefile index 3fd8236..32a5a7e 100644 --- a/Makefile +++ b/Makefile @@ -39,8 +39,8 @@ generate: mkdir go-tidy # Make server server: mkdir generate go-tidy libwhisper libggml - @echo "Building whisper-server" - @PKG_CONFIG_PATH=${ROOT_PATH}/${BUILD_DIR} ${GO} build ${BUILD_FLAGS} -o ${BUILD_DIR}/whisper-server ./cmd/server + @echo "Building whisper" + @PKG_CONFIG_PATH=${ROOT_PATH}/${BUILD_DIR} ${GO} build ${BUILD_FLAGS} -o ${BUILD_DIR}/whisper ./cmd/whisper # Make cli cli: mkdir generate go-tidy diff --git a/cmd/server/main.go b/cmd/server/main.go index cced3d4..e9fadac 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -13,7 +13,7 @@ import ( // Packages context "github.com/mutablelogic/go-server/pkg/context" httpserver "github.com/mutablelogic/go-server/pkg/httpserver" - whisper "github.com/mutablelogic/go-whisper/pkg/whisper" + whisper "github.com/mutablelogic/go-whisper" api "github.com/mutablelogic/go-whisper/pkg/whisper/api" ) diff --git a/cmd/whisper/download.go b/cmd/whisper/download.go new file mode 100644 index 0000000..ccfa6b6 --- /dev/null +++ b/cmd/whisper/download.go @@ -0,0 +1,22 @@ +package main + +import ( + "log" + + // Packages + "github.com/djthorpe/go-tablewriter" +) + +type DownloadCmd struct { + Model string `arg:"" help:"Model to download"` +} + +func (cmd *DownloadCmd) Run(ctx *Globals) error { + model, err := ctx.service.DownloadModel(ctx.ctx, cmd.Model, func(curBytes, totalBytes uint64) { + log.Printf("Downloaded %d of %d bytes", curBytes, totalBytes) + }) + if err != nil { + return err + } + return ctx.writer.Write(model, tablewriter.OptHeader()) +} diff --git a/cmd/whisper/main.go b/cmd/whisper/main.go new file mode 100644 index 0000000..c058800 --- /dev/null +++ b/cmd/whisper/main.go @@ -0,0 +1,101 @@ +package main + +import ( + "context" + "os" + "path/filepath" + "syscall" + + // Packages + kong "github.com/alecthomas/kong" + tablewriter "github.com/djthorpe/go-tablewriter" + ctx "github.com/mutablelogic/go-server/pkg/context" + whisper "github.com/mutablelogic/go-whisper" +) + +type Globals struct { + NoGPU bool `name:"nogpu" help:"Disable GPU acceleration"` + Debug bool `name:"debug" help:"Enable debug output"` + Dir string `name:"dir" help:"Path to model store, uses ${WHISPER_DIR} " default:"${WHISPER_DIR}"` + + // Writer, service and context + writer *tablewriter.Writer + service *whisper.Whisper + ctx context.Context +} + +type CLI struct { + Globals + Models ModelsCmd `cmd:"models" help:"List models"` + Download DownloadCmd `cmd:"download" help:"Download a model"` + Server ServerCmd `cmd:"models" help:"Run the whisper server"` +} + +func main() { + // The name of the executable + name, err := os.Executable() + if err != nil { + panic(err) + } else { + name = filepath.Base(name) + } + + // Create a cli parser + cli := CLI{} + cmd := kong.Parse(&cli, + kong.Name(name), + kong.Description("speech transcription and translation service"), + kong.UsageOnError(), + kong.ConfigureHelp(kong.HelpOptions{Compact: true}), + kong.Vars{ + "WHISPER_DIR": dirEnvOrDefault(name), + }, + ) + + // Create a whisper server - set options + opts := []whisper.Opt{} + if cli.Globals.Debug { + opts = append(opts, whisper.OptDebug()) + } + if cli.Globals.NoGPU { + opts = append(opts, whisper.OptNoGPU()) + } + + // Create directory if it doesn't exist + if err := os.MkdirAll(cli.Globals.Dir, 0755); err != nil { + cmd.FatalIfErrorf(err) + return + } + + // Create a whisper server - create + service, err := whisper.New(cli.Globals.Dir, opts...) + if err != nil { + cmd.FatalIfErrorf(err) + return + } else { + cli.Globals.service = service + } + defer service.Close() + + // Create a tablewriter object with text output + writer := tablewriter.New(os.Stdout, tablewriter.OptOutputText()) + cli.Globals.writer = writer + + // Create a context + cli.Globals.ctx = ctx.ContextForSignal(os.Interrupt, syscall.SIGQUIT) + + // Run the command + if err := cmd.Run(&cli.Globals); err != nil { + cmd.FatalIfErrorf(err) + } +} + +func dirEnvOrDefault(name string) string { + if dir := os.Getenv("WHISPER_DIR"); dir != "" { + return dir + } + if dir, err := os.UserCacheDir(); err == nil { + return filepath.Join(dir, name) + } + return os.TempDir() +} diff --git a/cmd/whisper/models.go b/cmd/whisper/models.go new file mode 100644 index 0000000..77549cb --- /dev/null +++ b/cmd/whisper/models.go @@ -0,0 +1,12 @@ +package main + +import ( + // Packages + "github.com/djthorpe/go-tablewriter" +) + +type ModelsCmd struct{} + +func (*ModelsCmd) Run(ctx *Globals) error { + return ctx.writer.Write(ctx.service.ListModels(), tablewriter.OptHeader()) +} diff --git a/cmd/whisper/server.go b/cmd/whisper/server.go new file mode 100644 index 0000000..bdce752 --- /dev/null +++ b/cmd/whisper/server.go @@ -0,0 +1,37 @@ +package main + +import ( + "log" + "net/http" + + // Packages + "github.com/mutablelogic/go-server/pkg/httpserver" + "github.com/mutablelogic/go-whisper/pkg/whisper/api" +) + +type ServerCmd struct { + Endpoint string `name:"endpoint" help:"Endpoint for the server" default:"/api/v1"` + Listen string `name:"listen" help:"Listen address for the server" default:"localhost:8080"` +} + +func (cmd *ServerCmd) Run(ctx *Globals) error { + // Create a mux for serving requests, then register the endpoints with the mux + mux := http.NewServeMux() + + // Register the endpoints + api.RegisterEndpoints(cmd.Endpoint, mux, ctx.service) + + // Create a new HTTP server + log.Println("List address", cmd.Listen) + server, err := httpserver.Config{ + Listen: cmd.Listen, + Router: mux, + }.New() + if err != nil { + return err + } + + // Run the server until CTRL+C + log.Println("Press CTRL+C to exit") + return server.Run(ctx.ctx) +} diff --git a/pkg/whisper/opt.go b/opt.go similarity index 100% rename from pkg/whisper/opt.go rename to opt.go diff --git a/pkg/whisper/api/models.go b/pkg/whisper/api/models.go index 89a515c..ebae153 100644 --- a/pkg/whisper/api/models.go +++ b/pkg/whisper/api/models.go @@ -11,7 +11,7 @@ import ( // Packages "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" "github.com/mutablelogic/go-whisper/pkg/whisper/schema" ) diff --git a/pkg/whisper/api/register.go b/pkg/whisper/api/register.go index 50ffd5e..153f3ec 100644 --- a/pkg/whisper/api/register.go +++ b/pkg/whisper/api/register.go @@ -6,7 +6,7 @@ import ( // Packages "github.com/mutablelogic/go-server/pkg/httpresponse" - "github.com/mutablelogic/go-whisper/pkg/whisper" + "github.com/mutablelogic/go-whisper" ) ///////////////////////////////////////////////////////////////////////////// diff --git a/pkg/whisper/api/transcribe.go b/pkg/whisper/api/transcribe.go index b2faef9..dddd549 100644 --- a/pkg/whisper/api/transcribe.go +++ b/pkg/whisper/api/transcribe.go @@ -10,7 +10,7 @@ import ( // Packages "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" "github.com/mutablelogic/go-whisper/pkg/whisper/schema" "github.com/mutablelogic/go-whisper/pkg/whisper/segmenter" "github.com/mutablelogic/go-whisper/pkg/whisper/task" diff --git a/pkg/whisper/pool/contextpool.go b/pkg/whisper/pool/contextpool.go index 6048aa7..15fe9a2 100644 --- a/pkg/whisper/pool/contextpool.go +++ b/pkg/whisper/pool/contextpool.go @@ -5,8 +5,7 @@ import ( "fmt" // Packages - - "github.com/mutablelogic/go-whisper/pkg/whisper/schema" + schema "github.com/mutablelogic/go-whisper/pkg/whisper/schema" task "github.com/mutablelogic/go-whisper/pkg/whisper/task" // Namespace imports diff --git a/pkg/whisper/whisper.go b/whisper.go similarity index 97% rename from pkg/whisper/whisper.go rename to whisper.go index b6357fb..f82604b 100644 --- a/pkg/whisper/whisper.go +++ b/whisper.go @@ -9,10 +9,10 @@ import ( "strings" // Packages - "github.com/mutablelogic/go-media/pkg/ffmpeg" + ffmpeg "github.com/mutablelogic/go-media/pkg/ffmpeg" model "github.com/mutablelogic/go-whisper/pkg/whisper/model" pool "github.com/mutablelogic/go-whisper/pkg/whisper/pool" - "github.com/mutablelogic/go-whisper/pkg/whisper/schema" + schema "github.com/mutablelogic/go-whisper/pkg/whisper/schema" task "github.com/mutablelogic/go-whisper/pkg/whisper/task" whisper "github.com/mutablelogic/go-whisper/sys/whisper" diff --git a/pkg/whisper/whisper_test.go b/whisper_test.go similarity index 97% rename from pkg/whisper/whisper_test.go rename to whisper_test.go index db8f304..1f8e4fc 100644 --- a/pkg/whisper/whisper_test.go +++ b/whisper_test.go @@ -8,7 +8,7 @@ import ( // Packages wav "github.com/go-audio/wav" - whisper "github.com/mutablelogic/go-whisper/pkg/whisper" + whisper "github.com/mutablelogic/go-whisper" task "github.com/mutablelogic/go-whisper/pkg/whisper/task" assert "github.com/stretchr/testify/assert" @@ -17,9 +17,9 @@ import ( ) const MODEL_TINY = "ggml-tiny.en-q5_1.bin" -const SAMPLE_EN = "../../samples/jfk.wav" -const SAMPLE_FR = "../../samples/OlivierL.wav" -const SAMPLE_DE = "../../samples/ge-podcast.wav" +const SAMPLE_EN = "samples/jfk.wav" +const SAMPLE_FR = "samples/OlivierL.wav" +const SAMPLE_DE = "samples/ge-podcast.wav" func Test_whisper_001(t *testing.T) { assert := assert.New(t)