go-whisper-api/whisper/model_pool.go

package whisper

import (
    "sync"

    "go-whisper-api/config"

    "github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
)

// Model is a loaded whisper.cpp weights file (re-export for API callers).
type Model = whisper.Model

// ModelPool keeps whisper models loaded in memory and serializes inference per model path.
type ModelPool struct {
    mu      sync.Mutex
    entries map[string]*pooledModel
}

type pooledModel struct {
    model whisper.Model
    mu    sync.Mutex
}

func NewModelPool() *ModelPool {
    return &ModelPool{entries: make(map[string]*pooledModel)}
}

func (p *ModelPool) WithModel(path string, fn func(whisper.Model) error) error {
    e, err := p.entry(path)
    if err != nil {
        return err
    }
    e.mu.Lock()
    defer e.mu.Unlock()
    return fn(e.model)
}

func (p *ModelPool) entry(path string) (*pooledModel, error) {
    p.mu.Lock()
    defer p.mu.Unlock()
    if e, ok := p.entries[path]; ok {
        return e, nil
    }
    m, err := whisper.New(path)
    if err != nil {
        return nil, err
    }
    e := &pooledModel{model: m}
    p.entries[path] = e
    return e, nil
}

func (p *ModelPool) Close() {
    p.mu.Lock()
    defer p.mu.Unlock()
    for _, e := range p.entries {
        _ = e.model.Close()
    }
    p.entries = make(map[string]*pooledModel)
}

var defaultPool = NewModelPool()

func DefaultPool() *ModelPool {
    return defaultPool
}

// Transcribe runs speech recognition using a cached model.
func Transcribe(cfg *config.Whisper) (TranscriptResult, error) {
    return TranscribeWithPool(defaultPool, cfg, RunOptions{})
}

func TranscribeWithPool(pool *ModelPool, cfg *config.Whisper, opts RunOptions) (TranscriptResult, error) {
    if pool == nil {
        pool = defaultPool
    }
    if err := cfg.Validate(); err != nil {
        return TranscriptResult{}, err
    }
    eng := &Engine{cfg: cfg, runOpts: opts}
    err := pool.WithModel(cfg.Model, func(m whisper.Model) error {
        return eng.transcribeWithModel(m)
    })
    if err != nil {
        return TranscriptResult{}, err
    }
    return eng.Result(), nil
}