* Added whisper state + default state on the whisper_context * Fixed some examples and bindings * Fixed whisper_n_len (which was used in some binding) and added whisper_n_len_from_state * Fixed comments * whisper : reuse kv_cache_free() and fix compiler warnings * whisper : clean-up the API comments --------- Co-authored-by: Sandro Hanea <sandrohanea@microsoft.com> Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
410 lines
14 KiB
Go
410 lines
14 KiB
Go
package whisper
|
|
|
|
import (
|
|
"errors"
|
|
"unsafe"
|
|
)
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
// CGO
|
|
|
|
/*
|
|
#cgo LDFLAGS: -lwhisper -lm -lstdc++
|
|
#cgo darwin LDFLAGS: -framework Accelerate
|
|
#include <whisper.h>
|
|
#include <stdlib.h>
|
|
|
|
extern void callNewSegment(void* user_data, int new);
|
|
extern bool callEncoderBegin(void* user_data);
|
|
|
|
// Text segment callback
|
|
// Called on every newly generated text segment
|
|
// Use the whisper_full_...() functions to obtain the text segments
|
|
static void whisper_new_segment_cb(struct whisper_context* ctx, struct whisper_state* state, int n_new, void* user_data) {
|
|
if(user_data != NULL && ctx != NULL) {
|
|
callNewSegment(user_data, n_new);
|
|
}
|
|
}
|
|
|
|
// Encoder begin callback
|
|
// If not NULL, called before the encoder starts
|
|
// If it returns false, the computation is aborted
|
|
static bool whisper_encoder_begin_cb(struct whisper_context* ctx, struct whisper_state* state, void* user_data) {
|
|
if(user_data != NULL && ctx != NULL) {
|
|
return callEncoderBegin(user_data);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
// Get default parameters and set callbacks
|
|
static struct whisper_full_params whisper_full_default_params_cb(struct whisper_context* ctx, enum whisper_sampling_strategy strategy) {
|
|
struct whisper_full_params params = whisper_full_default_params(strategy);
|
|
params.new_segment_callback = whisper_new_segment_cb;
|
|
params.new_segment_callback_user_data = (void*)(ctx);
|
|
params.encoder_begin_callback = whisper_encoder_begin_cb;
|
|
params.encoder_begin_callback_user_data = (void*)(ctx);
|
|
return params;
|
|
}
|
|
*/
|
|
import "C"
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
// TYPES
|
|
|
|
type (
|
|
Context C.struct_whisper_context
|
|
Token C.whisper_token
|
|
TokenData C.struct_whisper_token_data
|
|
SamplingStrategy C.enum_whisper_sampling_strategy
|
|
Params C.struct_whisper_full_params
|
|
)
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
// GLOBALS
|
|
|
|
const (
|
|
SAMPLING_GREEDY SamplingStrategy = C.WHISPER_SAMPLING_GREEDY
|
|
SAMPLING_BEAM_SEARCH SamplingStrategy = C.WHISPER_SAMPLING_BEAM_SEARCH
|
|
)
|
|
|
|
const (
|
|
SampleRate = C.WHISPER_SAMPLE_RATE // Expected sample rate, samples per second
|
|
SampleBits = uint16(unsafe.Sizeof(C.float(0))) * 8 // Sample size in bits
|
|
NumFFT = C.WHISPER_N_FFT
|
|
NumMEL = C.WHISPER_N_MEL
|
|
HopLength = C.WHISPER_HOP_LENGTH
|
|
ChunkSize = C.WHISPER_CHUNK_SIZE
|
|
)
|
|
|
|
var (
|
|
ErrTokenizerFailed = errors.New("whisper_tokenize failed")
|
|
ErrAutoDetectFailed = errors.New("whisper_lang_auto_detect failed")
|
|
ErrConversionFailed = errors.New("whisper_convert failed")
|
|
ErrInvalidLanguage = errors.New("invalid language")
|
|
)
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
// PUBLIC METHODS
|
|
|
|
// Allocates all memory needed for the model and loads the model from the given file.
|
|
// Returns NULL on failure.
|
|
func Whisper_init(path string) *Context {
|
|
cPath := C.CString(path)
|
|
defer C.free(unsafe.Pointer(cPath))
|
|
if ctx := C.whisper_init_from_file(cPath); ctx != nil {
|
|
return (*Context)(ctx)
|
|
} else {
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// Frees all memory allocated by the model.
|
|
func (ctx *Context) Whisper_free() {
|
|
C.whisper_free((*C.struct_whisper_context)(ctx))
|
|
}
|
|
|
|
// 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 {
|
|
if C.whisper_pcm_to_mel((*C.struct_whisper_context)(ctx), (*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
|
|
func (ctx *Context) Whisper_set_mel(data []float32, n_mel int) error {
|
|
if C.whisper_set_mel((*C.struct_whisper_context)(ctx), (*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.
|
|
func (ctx *Context) Whisper_encode(offset, threads int) error {
|
|
if C.whisper_encode((*C.struct_whisper_context)(ctx), 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.
|
|
// n_past is the number of tokens to use from previous decoder calls.
|
|
func (ctx *Context) Whisper_decode(tokens []Token, past, threads int) error {
|
|
if C.whisper_decode((*C.struct_whisper_context)(ctx), (*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) {
|
|
cText := C.CString(text)
|
|
defer C.free(unsafe.Pointer(cText))
|
|
if n := C.whisper_tokenize((*C.struct_whisper_context)(ctx), cText, (*C.whisper_token)(&tokens[0]), C.int(len(tokens))); n >= 0 {
|
|
return int(n), nil
|
|
} else {
|
|
return 0, ErrTokenizerFailed
|
|
}
|
|
}
|
|
|
|
// Return the id of the specified language, returns -1 if not found
|
|
// Examples:
|
|
//
|
|
// "de" -> 2
|
|
// "german" -> 2
|
|
func (ctx *Context) Whisper_lang_id(lang string) int {
|
|
return int(C.whisper_lang_id(C.CString(lang)))
|
|
}
|
|
|
|
// Largest language id (i.e. number of available languages - 1)
|
|
func Whisper_lang_max_id() int {
|
|
return int(C.whisper_lang_max_id())
|
|
}
|
|
|
|
// Return the short string of the specified language id (e.g. 2 -> "de"),
|
|
// returns empty string if not found
|
|
func Whisper_lang_str(id int) string {
|
|
return C.GoString(C.whisper_lang_str(C.int(id)))
|
|
}
|
|
|
|
// Use mel data at offset_ms to try and auto-detect the spoken language
|
|
// Make sure to call whisper_pcm_to_mel() or whisper_set_mel() first.
|
|
// Returns the probabilities of all languages.
|
|
// ref: https://github.com/openai/whisper/blob/main/whisper/decoding.py#L18-L69
|
|
func (ctx *Context) Whisper_lang_auto_detect(offset_ms, n_threads int) ([]float32, error) {
|
|
probs := make([]float32, Whisper_lang_max_id()+1)
|
|
if n := int(C.whisper_lang_auto_detect((*C.struct_whisper_context)(ctx), 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)))
|
|
}
|
|
|
|
func (ctx *Context) Whisper_n_vocab() int {
|
|
return int(C.whisper_n_vocab((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
func (ctx *Context) Whisper_n_text_ctx() int {
|
|
return int(C.whisper_n_text_ctx((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
func (ctx *Context) Whisper_n_audio_ctx() int {
|
|
return int(C.whisper_n_audio_ctx((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
func (ctx *Context) Whisper_is_multilingual() int {
|
|
return int(C.whisper_is_multilingual((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
// The probabilities for the next token
|
|
//func (ctx *Whisper_context) Whisper_get_probs() []float32 {
|
|
// return (*[1 << 30]float32)(unsafe.Pointer(C.whisper_get_probs((*C.struct_whisper_context)(ctx))))[:ctx.Whisper_n_vocab()]
|
|
//}
|
|
|
|
// Token Id -> String. Uses the vocabulary in the provided context
|
|
func (ctx *Context) Whisper_token_to_str(token Token) string {
|
|
return C.GoString(C.whisper_token_to_str((*C.struct_whisper_context)(ctx), C.whisper_token(token)))
|
|
}
|
|
|
|
// Special tokens
|
|
func (ctx *Context) Whisper_token_eot() Token {
|
|
return Token(C.whisper_token_eot((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
// Special tokens
|
|
func (ctx *Context) Whisper_token_sot() Token {
|
|
return Token(C.whisper_token_sot((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
// Special tokens
|
|
func (ctx *Context) Whisper_token_prev() Token {
|
|
return Token(C.whisper_token_prev((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
// Special tokens
|
|
func (ctx *Context) Whisper_token_solm() Token {
|
|
return Token(C.whisper_token_solm((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
// Special tokens
|
|
func (ctx *Context) Whisper_token_not() Token {
|
|
return Token(C.whisper_token_not((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
// Special tokens
|
|
func (ctx *Context) Whisper_token_beg() Token {
|
|
return Token(C.whisper_token_beg((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
// Special tokens
|
|
func (ctx *Context) Whisper_token_lang(lang_id int) Token {
|
|
return Token(C.whisper_token_lang((*C.struct_whisper_context)(ctx), C.int(lang_id)))
|
|
}
|
|
|
|
// Task tokens
|
|
func Whisper_token_translate() Token {
|
|
return Token(C.whisper_token_translate())
|
|
}
|
|
|
|
// Task tokens
|
|
func Whisper_token_transcribe() Token {
|
|
return Token(C.whisper_token_transcribe())
|
|
}
|
|
|
|
// Performance information
|
|
func (ctx *Context) Whisper_print_timings() {
|
|
C.whisper_print_timings((*C.struct_whisper_context)(ctx))
|
|
}
|
|
|
|
// Performance information
|
|
func (ctx *Context) Whisper_reset_timings() {
|
|
C.whisper_reset_timings((*C.struct_whisper_context)(ctx))
|
|
}
|
|
|
|
// Print system information
|
|
func Whisper_print_system_info() string {
|
|
return C.GoString(C.whisper_print_system_info())
|
|
}
|
|
|
|
// Return default parameters for a strategy
|
|
func (ctx *Context) Whisper_full_default_params(strategy SamplingStrategy) Params {
|
|
// Get default parameters
|
|
return Params(C.whisper_full_default_params_cb((*C.struct_whisper_context)(ctx), C.enum_whisper_sampling_strategy(strategy)))
|
|
}
|
|
|
|
// Run the entire model: PCM -> log mel spectrogram -> encoder -> decoder -> text
|
|
// Uses the specified decoding strategy to obtain the text.
|
|
func (ctx *Context) Whisper_full(params Params, samples []float32, encoderBeginCallback func() bool, newSegmentCallback func(int)) error {
|
|
registerEncoderBeginCallback(ctx, encoderBeginCallback)
|
|
registerNewSegmentCallback(ctx, newSegmentCallback)
|
|
defer registerEncoderBeginCallback(ctx, nil)
|
|
defer registerNewSegmentCallback(ctx, nil)
|
|
if C.whisper_full((*C.struct_whisper_context)(ctx), (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.
|
|
func (ctx *Context) Whisper_full_parallel(params Params, samples []float32, processors int, encoderBeginCallback func() bool, newSegmentCallback func(int)) error {
|
|
registerEncoderBeginCallback(ctx, encoderBeginCallback)
|
|
registerNewSegmentCallback(ctx, newSegmentCallback)
|
|
defer registerEncoderBeginCallback(ctx, nil)
|
|
defer registerNewSegmentCallback(ctx, nil)
|
|
|
|
if C.whisper_full_parallel((*C.struct_whisper_context)(ctx), (C.struct_whisper_full_params)(params), (*C.float)(&samples[0]), C.int(len(samples)), C.int(processors)) == 0 {
|
|
return nil
|
|
} else {
|
|
return ErrConversionFailed
|
|
}
|
|
}
|
|
|
|
// Number of generated text segments.
|
|
// A segment can be a few words, a sentence, or even a paragraph.
|
|
func (ctx *Context) Whisper_full_n_segments() int {
|
|
return int(C.whisper_full_n_segments((*C.struct_whisper_context)(ctx)))
|
|
}
|
|
|
|
// 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)))
|
|
}
|
|
|
|
// 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)))
|
|
}
|
|
|
|
// 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)))
|
|
}
|
|
|
|
// 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)))
|
|
}
|
|
|
|
// 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)))
|
|
}
|
|
|
|
// 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)))
|
|
}
|
|
|
|
// 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)))
|
|
}
|
|
|
|
// 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)))
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
// CALLBACKS
|
|
|
|
var (
|
|
cbNewSegment = make(map[unsafe.Pointer]func(int))
|
|
cbEncoderBegin = make(map[unsafe.Pointer]func() bool)
|
|
)
|
|
|
|
func registerNewSegmentCallback(ctx *Context, fn func(int)) {
|
|
if fn == nil {
|
|
delete(cbNewSegment, unsafe.Pointer(ctx))
|
|
} else {
|
|
cbNewSegment[unsafe.Pointer(ctx)] = fn
|
|
}
|
|
}
|
|
|
|
func registerEncoderBeginCallback(ctx *Context, fn func() bool) {
|
|
if fn == nil {
|
|
delete(cbEncoderBegin, unsafe.Pointer(ctx))
|
|
} else {
|
|
cbEncoderBegin[unsafe.Pointer(ctx)] = fn
|
|
}
|
|
}
|
|
|
|
//export callNewSegment
|
|
func callNewSegment(user_data unsafe.Pointer, new C.int) {
|
|
if fn, ok := cbNewSegment[user_data]; ok {
|
|
fn(int(new))
|
|
}
|
|
}
|
|
|
|
//export callEncoderBegin
|
|
func callEncoderBegin(user_data unsafe.Pointer) C.bool {
|
|
if fn, ok := cbEncoderBegin[user_data]; ok {
|
|
if fn() {
|
|
return C.bool(true)
|
|
} else {
|
|
return C.bool(false)
|
|
}
|
|
}
|
|
return true
|
|
}
|