kami
120 lines
2.9 KiB
Go
120 lines
2.9 KiB
Go
package main
|
|
|
|
/*
|
|
#cgo CFLAGS: -I${SRCDIR}/../../deps/include -I${SRCDIR}/../../deps/whisper.cpp/ggml/include
|
|
#cgo LDFLAGS: -L${SRCDIR}/../../deps/lib -Wl,-rpath,${SRCDIR}/../../deps/lib -lwhisper -lggml -lggml-base -lggml-cpu -lggml-vulkan -lm -lstdc++ -fopenmp
|
|
#include <whisper.h>
|
|
#include <stdlib.h>
|
|
*/
|
|
import "C"
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math"
|
|
"unsafe"
|
|
|
|
"github.com/kami/maven/internal/worker"
|
|
)
|
|
|
|
type whisperHandler struct {
|
|
ctx *C.struct_whisper_context
|
|
}
|
|
|
|
func newWhisperHandler(modelPath string) (*whisperHandler, error) {
|
|
cparams := C.whisper_context_default_params()
|
|
cPath := C.CString(modelPath)
|
|
defer C.free(unsafe.Pointer(cPath))
|
|
|
|
ctx := C.whisper_init_from_file_with_params(cPath, cparams)
|
|
if ctx == nil {
|
|
return nil, fmt.Errorf("whisper: failed to init from %s", modelPath)
|
|
}
|
|
return &whisperHandler{ctx: ctx}, nil
|
|
}
|
|
|
|
func (h *whisperHandler) Transcribe(ctx context.Context, req worker.TranscribeReq) (worker.TranscribeResp, error) {
|
|
a := req.Audio
|
|
if len(a.Bytes) == 0 {
|
|
return worker.TranscribeResp{}, fmt.Errorf("whisper: empty audio")
|
|
}
|
|
|
|
nSamples := len(a.Bytes) / 2
|
|
samples := make([]float32, nSamples)
|
|
for i := 0; i < nSamples; i++ {
|
|
s := int16(a.Bytes[i*2]) | int16(a.Bytes[i*2+1])<<8
|
|
samples[i] = float32(s) / 32768.0
|
|
}
|
|
|
|
params := C.whisper_full_default_params(C.WHISPER_SAMPLING_GREEDY)
|
|
params.print_progress = false
|
|
params.print_realtime = false
|
|
params.print_timestamps = false
|
|
params.print_special = false
|
|
params.n_threads = C.int(4)
|
|
params.single_segment = true
|
|
|
|
lang := C.CString(req.Lang)
|
|
defer C.free(unsafe.Pointer(lang))
|
|
params.language = lang
|
|
params.detect_language = false
|
|
|
|
cSamples := (*C.float)(unsafe.Pointer(&samples[0]))
|
|
res := C.whisper_full(h.ctx, params, cSamples, C.int(nSamples))
|
|
if res != 0 {
|
|
return worker.TranscribeResp{}, fmt.Errorf("whisper: full failed: %d", int(res))
|
|
}
|
|
|
|
nSegments := int(C.whisper_full_n_segments(h.ctx))
|
|
if nSegments == 0 {
|
|
return worker.TranscribeResp{Text: "", Confidence: 0}, nil
|
|
}
|
|
|
|
var text string
|
|
totalLogProb := float64(0)
|
|
totalTokens := 0
|
|
|
|
for i := 0; i < nSegments; i++ {
|
|
cSeg := C.whisper_full_get_segment_text(h.ctx, C.int(i))
|
|
if cSeg != nil {
|
|
text += C.GoString(cSeg)
|
|
}
|
|
|
|
nTokens := int(C.whisper_full_n_tokens(h.ctx, C.int(i)))
|
|
for j := 0; j < nTokens; j++ {
|
|
p := float64(C.whisper_full_get_token_p(h.ctx, C.int(i), C.int(j)))
|
|
if p > 0 {
|
|
totalLogProb += math.Log(p)
|
|
totalTokens++
|
|
}
|
|
}
|
|
}
|
|
|
|
confidence := 0.0
|
|
if totalTokens > 0 {
|
|
avgLogProb := totalLogProb / float64(totalTokens)
|
|
confidence = math.Exp(avgLogProb)
|
|
}
|
|
|
|
noSpeechProb := float64(C.whisper_full_get_segment_no_speech_prob(h.ctx, 0))
|
|
if noSpeechProb > 0.9 {
|
|
confidence = 0
|
|
}
|
|
|
|
if math.IsNaN(confidence) || math.IsInf(confidence, 0) {
|
|
confidence = 0
|
|
}
|
|
|
|
return worker.TranscribeResp{
|
|
Text: text,
|
|
Confidence: confidence,
|
|
}, nil
|
|
}
|
|
|
|
func (h *whisperHandler) Close() {
|
|
if h.ctx != nil {
|
|
C.whisper_free(h.ctx)
|
|
h.ctx = nil
|
|
}
|
|
}
|