From 59b3c4723cb10ad10c7d8a8db59e076d29bbe117 Mon Sep 17 00:00:00 2001 From: ciricc Date: Sat, 13 Sep 2025 15:31:31 +0300 Subject: [PATCH] feat(go bindings): add state abstraction --- bindings/go/pkg/whisper/consts.go | 1 + bindings/go/pkg/whisper/context.go | 122 ++++++++----------- bindings/go/pkg/whisper/interface.go | 9 ++ bindings/go/pkg/whisper/model.go | 16 +++ bindings/go/pkg/whisper/state.go | 125 +++++++++++++++++++ bindings/go/pkg/whisper/state_test.go | 128 ++++++++++++++++++++ bindings/go/whisper.go | 167 +++++++++++++++++++++++--- bindings/go/whisper_test.go | 160 +++++++++++++++++++++++- 8 files changed, 640 insertions(+), 88 deletions(-) create mode 100644 bindings/go/pkg/whisper/state.go create mode 100644 bindings/go/pkg/whisper/state_test.go diff --git a/bindings/go/pkg/whisper/consts.go b/bindings/go/pkg/whisper/consts.go index 5c22dc13a..6c778d3d2 100644 --- a/bindings/go/pkg/whisper/consts.go +++ b/bindings/go/pkg/whisper/consts.go @@ -16,6 +16,7 @@ var ( ErrProcessingFailed = errors.New("processing failed") ErrUnsupportedLanguage = errors.New("unsupported language") ErrModelNotMultilingual = errors.New("model is not multilingual") + ErrUnableToCreateState = errors.New("unable to create state") ) /////////////////////////////////////////////////////////////////////////////// diff --git a/bindings/go/pkg/whisper/context.go b/bindings/go/pkg/whisper/context.go index cb3d9eb8c..a7df02a66 100644 --- a/bindings/go/pkg/whisper/context.go +++ b/bindings/go/pkg/whisper/context.go @@ -20,9 +20,6 @@ type context struct { params whisper.Params } -// Make sure context adheres to the interface -var _ Context = (*context)(nil) - /////////////////////////////////////////////////////////////////////////////// // LIFECYCLE @@ -241,7 +238,7 @@ func (context *context) Process( return nil } -// Return the next segment of tokens +// NextSegment returns the next segment from the context buffer func (context *context) NextSegment() (Segment, error) { if context.model.ctx == nil { return Segment{}, ErrInternalAppError @@ -249,76 +246,11 @@ func (context *context) NextSegment() (Segment, error) { if context.n >= context.model.ctx.Whisper_full_n_segments() { return Segment{}, io.EOF } - - // Populate result result := toSegment(context.model.ctx, context.n) - - // Increment the cursor context.n++ - - // Return success return result, nil } -// Test for text tokens -func (context *context) IsText(t Token) bool { - switch { - case context.IsBEG(t): - return false - case context.IsSOT(t): - return false - case whisper.Token(t.Id) >= context.model.ctx.Whisper_token_eot(): - return false - case context.IsPREV(t): - return false - case context.IsSOLM(t): - return false - case context.IsNOT(t): - return false - default: - return true - } -} - -// Test for "begin" token -func (context *context) IsBEG(t Token) bool { - return whisper.Token(t.Id) == context.model.ctx.Whisper_token_beg() -} - -// Test for "start of transcription" token -func (context *context) IsSOT(t Token) bool { - return whisper.Token(t.Id) == context.model.ctx.Whisper_token_sot() -} - -// Test for "end of transcription" token -func (context *context) IsEOT(t Token) bool { - return whisper.Token(t.Id) == context.model.ctx.Whisper_token_eot() -} - -// Test for "start of prev" token -func (context *context) IsPREV(t Token) bool { - return whisper.Token(t.Id) == context.model.ctx.Whisper_token_prev() -} - -// Test for "start of lm" token -func (context *context) IsSOLM(t Token) bool { - return whisper.Token(t.Id) == context.model.ctx.Whisper_token_solm() -} - -// Test for "No timestamps" token -func (context *context) IsNOT(t Token) bool { - return whisper.Token(t.Id) == context.model.ctx.Whisper_token_not() -} - -// Test for token associated with a specific language -func (context *context) IsLANG(t Token, lang string) bool { - if id := context.model.ctx.Whisper_lang_id(lang); id >= 0 { - return whisper.Token(t.Id) == context.model.ctx.Whisper_token_lang(id) - } else { - return false - } -} - /////////////////////////////////////////////////////////////////////////////// // PRIVATE METHODS @@ -347,3 +279,55 @@ func toTokens(ctx *whisper.Context, n int) []Token { } return result } + +// Token helpers +func (context *context) IsText(t Token) bool { + switch { + case context.IsBEG(t): + return false + case context.IsSOT(t): + return false + case whisper.Token(t.Id) >= context.model.ctx.Whisper_token_eot(): + return false + case context.IsPREV(t): + return false + case context.IsSOLM(t): + return false + case context.IsNOT(t): + return false + default: + return true + } +} + +func (context *context) IsBEG(t Token) bool { + return whisper.Token(t.Id) == context.model.ctx.Whisper_token_beg() +} + +func (context *context) IsSOT(t Token) bool { + return whisper.Token(t.Id) == context.model.ctx.Whisper_token_sot() +} + +func (context *context) IsEOT(t Token) bool { + return whisper.Token(t.Id) == context.model.ctx.Whisper_token_eot() +} + +func (context *context) IsPREV(t Token) bool { + return whisper.Token(t.Id) == context.model.ctx.Whisper_token_prev() +} + +func (context *context) IsSOLM(t Token) bool { + return whisper.Token(t.Id) == context.model.ctx.Whisper_token_solm() +} + +func (context *context) IsNOT(t Token) bool { + return whisper.Token(t.Id) == context.model.ctx.Whisper_token_not() +} + +func (context *context) IsLANG(t Token, lang string) bool { + if id := context.model.ctx.Whisper_lang_id(lang); id >= 0 { + return whisper.Token(t.Id) == context.model.ctx.Whisper_token_lang(id) + } else { + return false + } +} diff --git a/bindings/go/pkg/whisper/interface.go b/bindings/go/pkg/whisper/interface.go index e3122c44b..b4705ec1c 100644 --- a/bindings/go/pkg/whisper/interface.go +++ b/bindings/go/pkg/whisper/interface.go @@ -85,6 +85,15 @@ type Context interface { SystemInfo() string } +// State is a per-request speech recognition state which shares the loaded model +// but isolates recognition results. It embeds Context, so any state-specific +// methods can be added later without breaking existing API. +type State interface { + io.Closer + + Context +} + // Segment is the text result of a speech recognition. type Segment struct { // Segment Number diff --git a/bindings/go/pkg/whisper/model.go b/bindings/go/pkg/whisper/model.go index 68a150223..0142a787a 100644 --- a/bindings/go/pkg/whisper/model.go +++ b/bindings/go/pkg/whisper/model.go @@ -99,3 +99,19 @@ func (model *model) NewContext() (Context, error) { // Return new context return newContext(model, params) } + +// NewState returns a new per-request state sharing the loaded model +func (model *model) NewState() (State, error) { + if model.ctx == nil { + return nil, ErrInternalAppError + } + params := model.ctx.Whisper_full_default_params(whisper.SAMPLING_GREEDY) + params.SetTranslate(false) + params.SetPrintSpecial(false) + params.SetPrintProgress(false) + params.SetPrintRealtime(false) + params.SetPrintTimestamps(false) + params.SetThreads(runtime.NumCPU()) + params.SetNoContext(true) + return newState(model, params) +} diff --git a/bindings/go/pkg/whisper/state.go b/bindings/go/pkg/whisper/state.go new file mode 100644 index 000000000..4ae81baf7 --- /dev/null +++ b/bindings/go/pkg/whisper/state.go @@ -0,0 +1,125 @@ +package whisper + +import ( + "io" + "strings" + "time" + + // Bindings + whisper "github.com/ggerganov/whisper.cpp/bindings/go" +) + +// state embeds context behavior and carries a low-level state pointer +// for isolated processing results. +type state struct { + *context + st *whisper.State +} + +// NewState creates a new per-request State from a Model without changing the Model interface. +func NewState(m Model) (State, error) { + impl, ok := m.(*model) + if !ok { + return nil, ErrInternalAppError + } + params := impl.ctx.Whisper_full_default_params(whisper.SAMPLING_GREEDY) + params.SetTranslate(false) + params.SetPrintSpecial(false) + params.SetPrintProgress(false) + params.SetPrintRealtime(false) + params.SetPrintTimestamps(false) + return newState(impl, params) +} + +// internal constructor used by model.NewState +func newState(model *model, params whisper.Params) (State, error) { + ctx := &context{model: model, params: params} + st := model.ctx.Whisper_init_state() + if st == nil { + return nil, ErrUnableToCreateState + } + return &state{context: ctx, st: st}, nil +} + +// Process using an isolated state for concurrency +func (s *state) Process( + data []float32, + callEncoderBegin EncoderBeginCallback, + callNewSegment SegmentCallback, + callProgress ProgressCallback, +) error { + if s.model.ctx == nil || s.st == nil { + return ErrInternalAppError + } + if callNewSegment != nil { + s.params.SetSingleSegment(true) + } + if err := s.model.ctx.Whisper_full_with_state(s.st, s.params, data, callEncoderBegin, + func(new int) { + if callNewSegment != nil { + num_segments := s.model.ctx.Whisper_full_n_segments_from_state(s.st) + s0 := num_segments - new + for i := s0; i < num_segments; i++ { + callNewSegment(toSegmentFromState(s.model.ctx, s.st, i)) + } + } + }, func(progress int) { + if callProgress != nil { + callProgress(progress) + } + }); err != nil { + return err + } + return nil +} + +// Return the next segment of tokens for state +func (s *state) NextSegment() (Segment, error) { + if s.model.ctx == nil { + return Segment{}, ErrInternalAppError + } + if s.n >= s.model.ctx.Whisper_full_n_segments_from_state(s.st) { + return Segment{}, io.EOF + } + result := toSegmentFromState(s.model.ctx, s.st, s.n) + s.n++ + return result, nil +} + +func (s *state) Close() error { + if s.st != nil { + s.st.Whisper_free_state() + s.st = nil + } + return nil +} + +// Helpers specific to state-based results +func toSegmentFromState(ctx *whisper.Context, st *whisper.State, n int) Segment { + return Segment{ + Num: n, + Text: stringsTrim(ctx.Whisper_full_get_segment_text_from_state(st, n)), + Start: duration10x(ctx.Whisper_full_get_segment_t0_from_state(st, n)), + End: duration10x(ctx.Whisper_full_get_segment_t1_from_state(st, n)), + Tokens: toTokensFromState(ctx, st, n), + } +} + +func toTokensFromState(ctx *whisper.Context, st *whisper.State, n int) []Token { + result := make([]Token, ctx.Whisper_full_n_tokens_from_state(st, n)) + for i := 0; i < len(result); i++ { + data := ctx.Whisper_full_get_token_data_from_state(st, n, i) + result[i] = Token{ + Id: int(ctx.Whisper_full_get_token_id_from_state(st, n, i)), + Text: ctx.Whisper_full_get_token_text_from_state(st, n, i), + P: ctx.Whisper_full_get_token_p_from_state(st, n, i), + Start: duration10x(data.T0()), + End: duration10x(data.T1()), + } + } + return result +} + +// small shared helpers to avoid importing time/strings here unnecessarily +func stringsTrim(s string) string { return strings.TrimSpace(s) } +func duration10x(ms10 int64) time.Duration { return time.Duration(ms10) * time.Millisecond * 10 } diff --git a/bindings/go/pkg/whisper/state_test.go b/bindings/go/pkg/whisper/state_test.go new file mode 100644 index 000000000..f893b8d2b --- /dev/null +++ b/bindings/go/pkg/whisper/state_test.go @@ -0,0 +1,128 @@ +package whisper_test + +import ( + "os" + "sync" + "testing" + + "github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper" + "github.com/go-audio/wav" + assert "github.com/stretchr/testify/assert" +) + +func TestState_Process(t *testing.T) { + assert := assert.New(t) + + if _, err := os.Stat(ModelPath); os.IsNotExist(err) { + t.Skip("Skipping test, model not found:", ModelPath) + } + if _, err := os.Stat(SamplePath); os.IsNotExist(err) { + t.Skip("Skipping test, sample not found:", SamplePath) + } + + fh, err := os.Open(SamplePath) + assert.NoError(err) + defer fh.Close() + + dec := wav.NewDecoder(fh) + buf, err := dec.FullPCMBuffer() + assert.NoError(err) + assert.Equal(uint16(1), dec.NumChans) + data := buf.AsFloat32Buffer().Data + + model, err := whisper.New(ModelPath) + assert.NoError(err) + assert.NotNil(model) + defer model.Close() + + st, err := whisper.NewState(model) + assert.NoError(err) + assert.NotNil(st) + defer func() { _ = st.Close() }() + + err = st.Process(data, nil, nil, nil) + assert.NoError(err) + + seg, err := st.NextSegment() + assert.NoError(err) + assert.NotEmpty(seg.Text) +} + +func TestState_Parallel_DifferentInputs(t *testing.T) { + assert := assert.New(t) + + if _, err := os.Stat(ModelPath); os.IsNotExist(err) { + t.Skip("Skipping test, model not found:", ModelPath) + } + if _, err := os.Stat(SamplePath); os.IsNotExist(err) { + t.Skip("Skipping test, sample not found:", SamplePath) + } + + fh, err := os.Open(SamplePath) + assert.NoError(err) + defer fh.Close() + + dec := wav.NewDecoder(fh) + buf, err := dec.FullPCMBuffer() + assert.NoError(err) + assert.Equal(uint16(1), dec.NumChans) + data := buf.AsFloat32Buffer().Data + assert.Greater(len(data), 10) + + // Create half-sample (second half) + half := make([]float32, len(data)/2) + copy(half, data[len(data)/2:]) + + model, err := whisper.New(ModelPath) + assert.NoError(err) + assert.NotNil(model) + defer model.Close() + + st1, err := whisper.NewState(model) + assert.NoError(err) + st2, err := whisper.NewState(model) + assert.NoError(err) + defer func() { _ = st1.Close() }() + defer func() { _ = st2.Close() }() + + // Run in parallel, but guard core call to respect context safety + var wg sync.WaitGroup + var first1, first2 string + var e1, e2 error + + wg.Add(2) + + // No mutex needed because each state is isolated + go func() { + defer wg.Done() + e1 = st1.Process(data, nil, nil, nil) + if e1 == nil { + seg, err := st1.NextSegment() + if err == nil { + first1 = seg.Text + } else { + e1 = err + } + } + }() + + go func() { + defer wg.Done() + e2 = st2.Process(half, nil, nil, nil) + if e2 == nil { + seg, err := st2.NextSegment() + if err == nil { + first2 = seg.Text + } else { + e2 = err + } + } + }() + + wg.Wait() + assert.NoError(e1) + assert.NoError(e2) + assert.NotEmpty(first1) + assert.NotEmpty(first2) + assert.NotEqual(first1, first2, "first segments should differ for different inputs") +} diff --git a/bindings/go/whisper.go b/bindings/go/whisper.go index 3ef73414d..cb5907ffe 100644 --- a/bindings/go/whisper.go +++ b/bindings/go/whisper.go @@ -2,6 +2,7 @@ package whisper import ( "errors" + "sync" "unsafe" ) @@ -67,6 +68,7 @@ import "C" type ( Context C.struct_whisper_context + State C.struct_whisper_state Token C.whisper_token TokenData C.struct_whisper_token_data SamplingStrategy C.enum_whisper_sampling_strategy @@ -116,6 +118,19 @@ func (ctx *Context) Whisper_free() { C.whisper_free((*C.struct_whisper_context)(ctx)) } +// Allocates a new state associated with the context. Returns nil on failure. +func (ctx *Context) Whisper_init_state() *State { + if s := C.whisper_init_state((*C.struct_whisper_context)(ctx)); s != nil { + return (*State)(s) + } + return nil +} + +// Frees all memory allocated by the state. +func (s *State) Whisper_free_state() { + C.whisper_free_state((*C.struct_whisper_state)(s)) +} + // Convert RAW PCM audio to log mel spectrogram. // The resulting spectrogram is stored inside the provided whisper context. func (ctx *Context) Whisper_pcm_to_mel(data []float32, threads int) error { @@ -126,6 +141,15 @@ func (ctx *Context) Whisper_pcm_to_mel(data []float32, threads int) error { } } +// Convert RAW PCM audio to log mel spectrogram into the provided state. +func (ctx *Context) Whisper_pcm_to_mel_with_state(state *State, data []float32, threads int) error { + if C.whisper_pcm_to_mel_with_state((*C.struct_whisper_context)(ctx), (*C.struct_whisper_state)(state), (*C.float)(&data[0]), C.int(len(data)), C.int(threads)) == 0 { + return nil + } else { + return ErrConversionFailed + } +} + // This can be used to set a custom log mel spectrogram inside the provided whisper context. // Use this instead of whisper_pcm_to_mel() if you want to provide your own log mel spectrogram. // n_mel must be 80 @@ -137,6 +161,15 @@ func (ctx *Context) Whisper_set_mel(data []float32, n_mel int) error { } } +// Set a custom log mel spectrogram into the provided state. +func (ctx *Context) Whisper_set_mel_with_state(state *State, data []float32, n_mel int) error { + if C.whisper_set_mel_with_state((*C.struct_whisper_context)(ctx), (*C.struct_whisper_state)(state), (*C.float)(&data[0]), C.int(len(data)), C.int(n_mel)) == 0 { + return nil + } else { + return ErrConversionFailed + } +} + // Run the Whisper encoder on the log mel spectrogram stored inside the provided whisper context. // Make sure to call whisper_pcm_to_mel() or whisper_set_mel() first. // offset can be used to specify the offset of the first frame in the spectrogram. @@ -148,6 +181,15 @@ func (ctx *Context) Whisper_encode(offset, threads int) error { } } +// Run the Whisper encoder using the provided state. +func (ctx *Context) Whisper_encode_with_state(state *State, offset, threads int) error { + if C.whisper_encode_with_state((*C.struct_whisper_context)(ctx), (*C.struct_whisper_state)(state), C.int(offset), C.int(threads)) == 0 { + return nil + } else { + return ErrConversionFailed + } +} + // Run the Whisper decoder to obtain the logits and probabilities for the next token. // Make sure to call whisper_encode() first. // tokens + n_tokens is the provided context for the decoder. @@ -160,6 +202,15 @@ func (ctx *Context) Whisper_decode(tokens []Token, past, threads int) error { } } +// Run the Whisper decoder using the provided state. +func (ctx *Context) Whisper_decode_with_state(state *State, tokens []Token, past, threads int) error { + if C.whisper_decode_with_state((*C.struct_whisper_context)(ctx), (*C.struct_whisper_state)(state), (*C.whisper_token)(&tokens[0]), C.int(len(tokens)), C.int(past), C.int(threads)) == 0 { + return nil + } else { + return ErrConversionFailed + } +} + // Convert the provided text into tokens. The tokens pointer must be large enough to hold the resulting tokens. // Returns the number of tokens on success func (ctx *Context) Whisper_tokenize(text string, tokens []Token) (int, error) { @@ -205,6 +256,16 @@ func (ctx *Context) Whisper_lang_auto_detect(offset_ms, n_threads int) ([]float3 } } +// Use mel data at offset_ms to auto-detect language using the provided state. +func (ctx *Context) Whisper_lang_auto_detect_with_state(state *State, offset_ms, n_threads int) ([]float32, error) { + probs := make([]float32, Whisper_lang_max_id()+1) + if n := int(C.whisper_lang_auto_detect_with_state((*C.struct_whisper_context)(ctx), (*C.struct_whisper_state)(state), C.int(offset_ms), C.int(n_threads), (*C.float)(&probs[0]))); n < 0 { + return nil, ErrAutoDetectFailed + } else { + return probs, nil + } +} + func (ctx *Context) Whisper_n_len() int { return int(C.whisper_n_len((*C.struct_whisper_context)(ctx))) } @@ -323,6 +384,28 @@ func (ctx *Context) Whisper_full( } } +// Run the entire model using the provided state: PCM -> mel -> encoder -> decoder -> text +func (ctx *Context) Whisper_full_with_state( + state *State, + params Params, + samples []float32, + encoderBeginCallback func() bool, + newSegmentCallback func(int), + progressCallback func(int), +) error { + registerEncoderBeginCallback(ctx, encoderBeginCallback) + registerNewSegmentCallback(ctx, newSegmentCallback) + registerProgressCallback(ctx, progressCallback) + defer registerEncoderBeginCallback(ctx, nil) + defer registerNewSegmentCallback(ctx, nil) + defer registerProgressCallback(ctx, nil) + if C.whisper_full_with_state((*C.struct_whisper_context)(ctx), (*C.struct_whisper_state)(state), (C.struct_whisper_full_params)(params), (*C.float)(&samples[0]), C.int(len(samples))) == 0 { + return nil + } else { + return ErrConversionFailed + } +} + // Split the input audio in chunks and process each chunk separately using whisper_full() // It seems this approach can offer some speedup in some cases. // However, the transcription accuracy can be worse at the beginning and end of each chunk. @@ -357,102 +440,152 @@ func (ctx *Context) Whisper_full_n_segments() int { return int(C.whisper_full_n_segments((*C.struct_whisper_context)(ctx))) } +func (ctx *Context) Whisper_full_n_segments_from_state(state *State) int { + return int(C.whisper_full_n_segments_from_state((*C.struct_whisper_state)(state))) +} + // Get the start and end time of the specified segment. func (ctx *Context) Whisper_full_get_segment_t0(segment int) int64 { return int64(C.whisper_full_get_segment_t0((*C.struct_whisper_context)(ctx), C.int(segment))) } +func (ctx *Context) Whisper_full_get_segment_t0_from_state(state *State, segment int) int64 { + return int64(C.whisper_full_get_segment_t0_from_state((*C.struct_whisper_state)(state), C.int(segment))) +} + // Get the start and end time of the specified segment. func (ctx *Context) Whisper_full_get_segment_t1(segment int) int64 { return int64(C.whisper_full_get_segment_t1((*C.struct_whisper_context)(ctx), C.int(segment))) } +func (ctx *Context) Whisper_full_get_segment_t1_from_state(state *State, segment int) int64 { + return int64(C.whisper_full_get_segment_t1_from_state((*C.struct_whisper_state)(state), C.int(segment))) +} + // Get the text of the specified segment. func (ctx *Context) Whisper_full_get_segment_text(segment int) string { return C.GoString(C.whisper_full_get_segment_text((*C.struct_whisper_context)(ctx), C.int(segment))) } +func (ctx *Context) Whisper_full_get_segment_text_from_state(state *State, segment int) string { + return C.GoString(C.whisper_full_get_segment_text_from_state((*C.struct_whisper_state)(state), C.int(segment))) +} + // Get number of tokens in the specified segment. func (ctx *Context) Whisper_full_n_tokens(segment int) int { return int(C.whisper_full_n_tokens((*C.struct_whisper_context)(ctx), C.int(segment))) } +func (ctx *Context) Whisper_full_n_tokens_from_state(state *State, segment int) int { + return int(C.whisper_full_n_tokens_from_state((*C.struct_whisper_state)(state), C.int(segment))) +} + // Get the token text of the specified token index in the specified segment. func (ctx *Context) Whisper_full_get_token_text(segment int, token int) string { return C.GoString(C.whisper_full_get_token_text((*C.struct_whisper_context)(ctx), C.int(segment), C.int(token))) } +func (ctx *Context) Whisper_full_get_token_text_from_state(state *State, segment int, token int) string { + return C.GoString(C.whisper_full_get_token_text_from_state((*C.struct_whisper_context)(ctx), (*C.struct_whisper_state)(state), C.int(segment), C.int(token))) +} + // Get the token of the specified token index in the specified segment. func (ctx *Context) Whisper_full_get_token_id(segment int, token int) Token { return Token(C.whisper_full_get_token_id((*C.struct_whisper_context)(ctx), C.int(segment), C.int(token))) } +func (ctx *Context) Whisper_full_get_token_id_from_state(state *State, segment int, token int) Token { + return Token(C.whisper_full_get_token_id_from_state((*C.struct_whisper_state)(state), C.int(segment), C.int(token))) +} + // Get token data for the specified token in the specified segment. // This contains probabilities, timestamps, etc. func (ctx *Context) Whisper_full_get_token_data(segment int, token int) TokenData { return TokenData(C.whisper_full_get_token_data((*C.struct_whisper_context)(ctx), C.int(segment), C.int(token))) } +func (ctx *Context) Whisper_full_get_token_data_from_state(state *State, segment int, token int) TokenData { + return TokenData(C.whisper_full_get_token_data_from_state((*C.struct_whisper_state)(state), C.int(segment), C.int(token))) +} + // Get the probability of the specified token in the specified segment. func (ctx *Context) Whisper_full_get_token_p(segment int, token int) float32 { return float32(C.whisper_full_get_token_p((*C.struct_whisper_context)(ctx), C.int(segment), C.int(token))) } +func (ctx *Context) Whisper_full_get_token_p_from_state(state *State, segment int, token int) float32 { + return float32(C.whisper_full_get_token_p_from_state((*C.struct_whisper_state)(state), C.int(segment), C.int(token))) +} + +func (ctx *Context) Whisper_full_lang_id_from_state(state *State) int { + return int(C.whisper_full_lang_id_from_state((*C.struct_whisper_state)(state))) +} + +func (ctx *Context) Whisper_n_len_from_state(state *State) int { + return int(C.whisper_n_len_from_state((*C.struct_whisper_state)(state))) +} + +func (ctx *Context) Whisper_get_logits_from_state(state *State) []float32 { + return (*[1 << 30]float32)(unsafe.Pointer(C.whisper_get_logits_from_state((*C.struct_whisper_state)(state))))[:ctx.Whisper_n_vocab()] +} + /////////////////////////////////////////////////////////////////////////////// // CALLBACKS var ( - cbNewSegment = make(map[unsafe.Pointer]func(int)) - cbProgress = make(map[unsafe.Pointer]func(int)) - cbEncoderBegin = make(map[unsafe.Pointer]func() bool) + cbNewSegment sync.Map // map[unsafe.Pointer]func(int) + cbProgress sync.Map // map[unsafe.Pointer]func(int) + cbEncoderBegin sync.Map // map[unsafe.Pointer]func() bool ) func registerNewSegmentCallback(ctx *Context, fn func(int)) { + k := unsafe.Pointer(ctx) if fn == nil { - delete(cbNewSegment, unsafe.Pointer(ctx)) + cbNewSegment.Delete(k) } else { - cbNewSegment[unsafe.Pointer(ctx)] = fn + cbNewSegment.Store(k, fn) } } func registerProgressCallback(ctx *Context, fn func(int)) { + k := unsafe.Pointer(ctx) if fn == nil { - delete(cbProgress, unsafe.Pointer(ctx)) + cbProgress.Delete(k) } else { - cbProgress[unsafe.Pointer(ctx)] = fn + cbProgress.Store(k, fn) } } func registerEncoderBeginCallback(ctx *Context, fn func() bool) { + k := unsafe.Pointer(ctx) if fn == nil { - delete(cbEncoderBegin, unsafe.Pointer(ctx)) + cbEncoderBegin.Delete(k) } else { - cbEncoderBegin[unsafe.Pointer(ctx)] = fn + cbEncoderBegin.Store(k, fn) } } //export callNewSegment func callNewSegment(user_data unsafe.Pointer, new C.int) { - if fn, ok := cbNewSegment[user_data]; ok { - fn(int(new)) + if v, ok := cbNewSegment.Load(user_data); ok { + v.(func(int))(int(new)) } } //export callProgress func callProgress(user_data unsafe.Pointer, progress C.int) { - if fn, ok := cbProgress[user_data]; ok { - fn(int(progress)) + if v, ok := cbProgress.Load(user_data); ok { + v.(func(int))(int(progress)) } } //export callEncoderBegin func callEncoderBegin(user_data unsafe.Pointer) C.bool { - if fn, ok := cbEncoderBegin[user_data]; ok { - if fn() { + if v, ok := cbEncoderBegin.Load(user_data); ok { + if v.(func() bool)() { return C.bool(true) - } else { - return C.bool(false) } + return C.bool(false) } return true } diff --git a/bindings/go/whisper_test.go b/bindings/go/whisper_test.go index 40648ffa8..23bbfbff0 100644 --- a/bindings/go/whisper_test.go +++ b/bindings/go/whisper_test.go @@ -1,8 +1,10 @@ package whisper_test import ( + "errors" "os" "runtime" + "sync" "testing" "time" @@ -39,7 +41,7 @@ func Test_Whisper_001(t *testing.T) { // Open samples fh, err := os.Open(SamplePath) assert.NoError(err) - defer fh.Close() + defer func() { _ = fh.Close() }() // Read samples d := wav.NewDecoder(fh) @@ -89,7 +91,7 @@ func Test_Whisper_003(t *testing.T) { // Open samples fh, err := os.Open(SamplePath) assert.NoError(err) - defer fh.Close() + defer func() { _ = fh.Close() }() // Read samples d := wav.NewDecoder(fh) @@ -111,3 +113,157 @@ func Test_Whisper_003(t *testing.T) { t.Logf("%s: %f", whisper.Whisper_lang_str(i), p) } } + +func Test_Whisper_State_Init_Free(t *testing.T) { + assert := assert.New(t) + if _, err := os.Stat(ModelPath); os.IsNotExist(err) { + t.Skip("Skipping test, model not found:", ModelPath) + } + + ctx := whisper.Whisper_init(ModelPath) + assert.NotNil(ctx) + defer ctx.Whisper_free() + + state := ctx.Whisper_init_state() + assert.NotNil(state) + state.Whisper_free_state() +} + +func Test_Whisper_Full_With_State(t *testing.T) { + assert := assert.New(t) + if _, err := os.Stat(ModelPath); os.IsNotExist(err) { + t.Skip("Skipping test, model not found:", ModelPath) + } + if _, err := os.Stat(SamplePath); os.IsNotExist(err) { + t.Skip("Skipping test, sample not found:", SamplePath) + } + + // Open samples + fh, err := os.Open(SamplePath) + assert.NoError(err) + defer func() { _ = fh.Close() }() + + // Read samples + d := wav.NewDecoder(fh) + buf, err := d.FullPCMBuffer() + assert.NoError(err) + data := buf.AsFloat32Buffer().Data + + ctx := whisper.Whisper_init(ModelPath) + assert.NotNil(ctx) + defer ctx.Whisper_free() + + state := ctx.Whisper_init_state() + assert.NotNil(state) + defer state.Whisper_free_state() + + params := ctx.Whisper_full_default_params(whisper.SAMPLING_GREEDY) + // Run using state + err = ctx.Whisper_full_with_state(state, params, data, nil, nil, nil) + assert.NoError(err) + + // Validate results are stored in state + nSegments := ctx.Whisper_full_n_segments_from_state(state) + assert.GreaterOrEqual(nSegments, 1) + text := ctx.Whisper_full_get_segment_text_from_state(state, 0) + assert.NotEmpty(text) +} + +func Test_Whisper_Lang_Auto_Detect_With_State(t *testing.T) { + assert := assert.New(t) + if _, err := os.Stat(ModelPath); os.IsNotExist(err) { + t.Skip("Skipping test, model not found:", ModelPath) + } + if _, err := os.Stat(SamplePath); os.IsNotExist(err) { + t.Skip("Skipping test, sample not found:", SamplePath) + } + + // Open samples + fh, err := os.Open(SamplePath) + assert.NoError(err) + defer func() { _ = fh.Close() }() + + // Read samples + d := wav.NewDecoder(fh) + buf, err := d.FullPCMBuffer() + assert.NoError(err) + data := buf.AsFloat32Buffer().Data + + ctx := whisper.Whisper_init(ModelPath) + assert.NotNil(ctx) + defer ctx.Whisper_free() + + state := ctx.Whisper_init_state() + assert.NotNil(state) + defer state.Whisper_free_state() + + threads := runtime.NumCPU() + // Prepare mel into state then detect + assert.NoError(ctx.Whisper_pcm_to_mel_with_state(state, data, threads)) + probs, err := ctx.Whisper_lang_auto_detect_with_state(state, 0, threads) + assert.NoError(err) + assert.Equal(whisper.Whisper_lang_max_id()+1, len(probs)) +} + +func Test_Whisper_Concurrent_With_State(t *testing.T) { + assert := assert.New(t) + if _, err := os.Stat(ModelPath); os.IsNotExist(err) { + t.Skip("Skipping test, model not found:", ModelPath) + } + if _, err := os.Stat(SamplePath); os.IsNotExist(err) { + t.Skip("Skipping test, sample not found:", SamplePath) + } + + // Load audio once + fh, err := os.Open(SamplePath) + assert.NoError(err) + defer func() { _ = fh.Close() }() + dec := wav.NewDecoder(fh) + buf, err := dec.FullPCMBuffer() + assert.NoError(err) + data := buf.AsFloat32Buffer().Data + + ctx := whisper.Whisper_init(ModelPath) + assert.NotNil(ctx) + defer ctx.Whisper_free() + + // Each goroutine has its own state + state1 := ctx.Whisper_init_state() + state2 := ctx.Whisper_init_state() + assert.NotNil(state1) + assert.NotNil(state2) + defer state1.Whisper_free_state() + defer state2.Whisper_free_state() + + params := ctx.Whisper_full_default_params(whisper.SAMPLING_GREEDY) + + var wg sync.WaitGroup + var mu sync.Mutex // guard calls into shared ctx, per upstream note not thread-safe for same context + errs := make(chan error, 2) + + worker := func(state *whisper.State) { + defer wg.Done() + mu.Lock() + err := ctx.Whisper_full_with_state(state, params, data, nil, nil, nil) + if err == nil { + n := ctx.Whisper_full_n_segments_from_state(state) + if n <= 0 { + err = errors.New("no segments") + } else { + _ = ctx.Whisper_full_get_segment_text_from_state(state, 0) + } + } + mu.Unlock() + errs <- err + } + + wg.Add(2) + go worker(state1) + go worker(state2) + wg.Wait() + close(errs) + + for e := range errs { + assert.NoError(e) + } +}