first commit
This commit is contained in:
@@ -0,0 +1,11 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"go-whisper-api/transcode"
|
||||
)
|
||||
|
||||
func AudioToWav(src, dst string) error {
|
||||
return transcode.ToWhisperWAV(context.Background(), src, dst)
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"go-whisper-api/config"
|
||||
|
||||
"github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
|
||||
"github.com/go-audio/wav"
|
||||
)
|
||||
|
||||
// LoadPCM16Mono reads 16 kHz mono WAV into float32 samples (for diarization).
|
||||
func LoadPCM16Mono(path string) ([]float32, error) {
|
||||
return loadPCM16Mono(path)
|
||||
}
|
||||
|
||||
func loadPCM16Mono(path string) ([]float32, error) {
|
||||
fh, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer fh.Close()
|
||||
dec := wav.NewDecoder(fh)
|
||||
buf, err := dec.FullPCMBuffer()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if dec.SampleRate != whisper.SampleRate {
|
||||
return nil, fmt.Errorf("unsupported sample rate: %d", dec.SampleRate)
|
||||
}
|
||||
if dec.NumChans != 1 {
|
||||
return nil, fmt.Errorf("unsupported number of channels: %d", dec.NumChans)
|
||||
}
|
||||
return buf.AsFloat32Buffer().Data, nil
|
||||
}
|
||||
|
||||
func prepareAudioPCM(sourcePath string) (data []float32, cleanup func(), err error) {
|
||||
cleanup = func() {}
|
||||
if data, err = loadPCM16Mono(sourcePath); err == nil {
|
||||
return data, cleanup, nil
|
||||
}
|
||||
dir, err := config.MkdirTemp("go-whisper-api-whisper-*")
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
cleanup = func() { os.RemoveAll(dir) }
|
||||
converted := filepath.Join(dir, "converted.wav")
|
||||
if err := AudioToWav(sourcePath, converted); err != nil {
|
||||
cleanup()
|
||||
return nil, nil, err
|
||||
}
|
||||
data, err = loadPCM16Mono(converted)
|
||||
if err != nil {
|
||||
cleanup()
|
||||
return nil, nil, err
|
||||
}
|
||||
return data, cleanup, nil
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
wpkg "github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
|
||||
)
|
||||
|
||||
// Turn is a speaker-active time range from diarization (seconds).
|
||||
type Turn struct {
|
||||
Start float32
|
||||
End float32
|
||||
Speaker int
|
||||
}
|
||||
|
||||
// FormatOptions controls joining Whisper segments into one string with newlines.
|
||||
type FormatOptions struct {
|
||||
PauseGap time.Duration
|
||||
SpeakerLabel string
|
||||
UseSpeakers bool
|
||||
}
|
||||
|
||||
// FormatSegments joins segment texts with \n on long pauses and optional speaker labels.
|
||||
func FormatSegments(segments []wpkg.Segment, turns []Turn, opts FormatOptions) string {
|
||||
if len(segments) == 0 {
|
||||
return ""
|
||||
}
|
||||
if opts.PauseGap <= 0 {
|
||||
opts.PauseGap = 1500 * time.Millisecond
|
||||
}
|
||||
label := strings.TrimSpace(opts.SpeakerLabel)
|
||||
if label == "" {
|
||||
label = "Спикер"
|
||||
}
|
||||
|
||||
lines := make([]segmentLine, len(segments))
|
||||
for i, seg := range segments {
|
||||
lines[i] = segmentLine{
|
||||
Text: strings.TrimSpace(seg.Text),
|
||||
Start: seg.Start,
|
||||
End: seg.End,
|
||||
Speaker: -1,
|
||||
}
|
||||
}
|
||||
if opts.UseSpeakers && len(turns) > 0 {
|
||||
assignSpeakers(lines, turns)
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
prevSpeaker := -2
|
||||
prevIdx := -1
|
||||
for i, line := range lines {
|
||||
if line.Text == "" {
|
||||
continue
|
||||
}
|
||||
if b.Len() > 0 && prevIdx >= 0 {
|
||||
speakerBreak := opts.UseSpeakers && line.Speaker >= 0 && line.Speaker != prevSpeaker
|
||||
pauseBreak := line.Start-lines[prevIdx].End >= opts.PauseGap
|
||||
switch {
|
||||
case speakerBreak:
|
||||
b.WriteString("\n\n")
|
||||
fmt.Fprintf(&b, "%s %d: ", label, line.Speaker+1)
|
||||
case pauseBreak:
|
||||
b.WriteString("\n")
|
||||
default:
|
||||
if !strings.HasSuffix(b.String(), " ") && !strings.HasSuffix(b.String(), "\n") {
|
||||
b.WriteByte(' ')
|
||||
}
|
||||
}
|
||||
} else if opts.UseSpeakers && line.Speaker >= 0 {
|
||||
fmt.Fprintf(&b, "%s %d: ", label, line.Speaker+1)
|
||||
}
|
||||
b.WriteString(line.Text)
|
||||
if line.Speaker >= 0 {
|
||||
prevSpeaker = line.Speaker
|
||||
}
|
||||
prevIdx = i
|
||||
}
|
||||
return strings.TrimSpace(b.String())
|
||||
}
|
||||
|
||||
type segmentLine struct {
|
||||
Text string
|
||||
Start time.Duration
|
||||
End time.Duration
|
||||
Speaker int
|
||||
}
|
||||
|
||||
func assignSpeakers(lines []segmentLine, turns []Turn) {
|
||||
for i := range lines {
|
||||
mid := lines[i].Start + (lines[i].End-lines[i].Start)/2
|
||||
lines[i].Speaker = speakerAt(mid, turns)
|
||||
}
|
||||
}
|
||||
|
||||
func speakerAt(t time.Duration, turns []Turn) int {
|
||||
sec := float32(t.Seconds())
|
||||
bestSpeaker := -1
|
||||
bestOverlap := float32(0)
|
||||
for _, tr := range turns {
|
||||
if sec >= tr.Start && sec < tr.End {
|
||||
return tr.Speaker
|
||||
}
|
||||
overlap := intervalOverlap(sec, sec, tr.Start, tr.End)
|
||||
if overlap > bestOverlap {
|
||||
bestOverlap = overlap
|
||||
bestSpeaker = tr.Speaker
|
||||
}
|
||||
}
|
||||
return bestSpeaker
|
||||
}
|
||||
|
||||
func intervalOverlap(a0, a1, b0, b1 float32) float32 {
|
||||
start := max32(a0, b0)
|
||||
end := min32(a1, b1)
|
||||
if end <= start {
|
||||
return 0
|
||||
}
|
||||
return end - start
|
||||
}
|
||||
|
||||
func max32(a, b float32) float32 {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func min32(a, b float32) float32 {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// PunctuateSegments runs punctuation on each segment separately (preserves line breaks).
|
||||
func PunctuateSegments(segments []wpkg.Segment, restore func(text string) (string, error)) ([]wpkg.Segment, error) {
|
||||
out := make([]wpkg.Segment, len(segments))
|
||||
copy(out, segments)
|
||||
for i := range out {
|
||||
t := strings.TrimSpace(out[i].Text)
|
||||
if t == "" {
|
||||
continue
|
||||
}
|
||||
p, err := restore(t)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[i].Text = " " + strings.TrimSpace(p)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
wpkg "github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
|
||||
)
|
||||
|
||||
func TestFormatSegments_pauseAndSpeaker(t *testing.T) {
|
||||
segments := []wpkg.Segment{
|
||||
{Text: " привет", Start: 0, End: 2 * time.Second},
|
||||
{Text: " мир", Start: 4 * time.Second, End: 5 * time.Second},
|
||||
{Text: " ответ", Start: 6 * time.Second, End: 8 * time.Second},
|
||||
}
|
||||
turns := []Turn{
|
||||
{Start: 0, End: 5.5, Speaker: 0},
|
||||
{Start: 5.5, End: 10, Speaker: 1},
|
||||
}
|
||||
got := FormatSegments(segments, turns, FormatOptions{
|
||||
PauseGap: 1500 * time.Millisecond,
|
||||
SpeakerLabel: "Спикер",
|
||||
UseSpeakers: true,
|
||||
})
|
||||
want := "Спикер 1: привет\nмир\n\nСпикер 2: ответ"
|
||||
if got != want {
|
||||
t.Fatalf("got %q\nwant %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatSegments_noSpeakers(t *testing.T) {
|
||||
segments := []wpkg.Segment{
|
||||
{Text: "a", Start: 0, End: time.Second},
|
||||
{Text: "b", Start: 3 * time.Second, End: 4 * time.Second},
|
||||
}
|
||||
got := FormatSegments(segments, nil, FormatOptions{PauseGap: time.Second, UseSpeakers: false})
|
||||
if got != "a\nb" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
func srtTimestamp(t time.Duration) string {
|
||||
return fmt.Sprintf("%02d:%02d:%02d,%03d",
|
||||
t/time.Hour,
|
||||
(t%time.Hour)/time.Minute,
|
||||
(t%time.Minute)/time.Second,
|
||||
(t%time.Second)/time.Millisecond,
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSrtTimestamp(t *testing.T) {
|
||||
type args struct {
|
||||
t time.Duration
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
args args
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "test 1",
|
||||
args: args{
|
||||
t: time.Duration(1*time.Hour + 2*time.Minute + 3*time.Second + 4*time.Millisecond),
|
||||
},
|
||||
want: "01:02:03,004",
|
||||
},
|
||||
{
|
||||
name: "test 2",
|
||||
args: args{
|
||||
t: time.Duration(10*time.Hour + 20*time.Minute + 30*time.Second + 40*time.Millisecond),
|
||||
},
|
||||
want: "10:20:30,040",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := srtTimestamp(tt.args.t); got != tt.want {
|
||||
t.Errorf("srtTimestamp() = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package whisper
|
||||
|
||||
// RunOptions affects transcription output formatting and optional diarization hints.
|
||||
type RunOptions struct {
|
||||
Format FormatOptions
|
||||
Turns []Turn
|
||||
PunctuateRestore func(text string) (string, error)
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
"go-whisper-api/config"
|
||||
|
||||
pkg "github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
|
||||
)
|
||||
|
||||
func ApplyVAD(ctx pkg.Context, vad config.VAD) {
|
||||
if !vad.Enabled {
|
||||
return
|
||||
}
|
||||
vad = vad.WithDefaults()
|
||||
ctx.SetVAD(true)
|
||||
ctx.SetVADModelPath(vad.Model)
|
||||
ctx.SetVADThreshold(float32(vad.Threshold))
|
||||
ctx.SetVADMinSpeechMs(vad.MinSpeechMs)
|
||||
ctx.SetVADMinSilenceMs(vad.MinSilenceMs)
|
||||
if vad.MaxSpeechSec > 0 {
|
||||
ctx.SetVADMaxSpeechSec(float32(vad.MaxSpeechSec))
|
||||
} else {
|
||||
ctx.SetVADMaxSpeechSec(float32(math.MaxFloat32))
|
||||
}
|
||||
ctx.SetVADSpeechPadMs(vad.SpeechPadMs)
|
||||
ctx.SetVADSamplesOverlap(float32(vad.SamplesOverlap))
|
||||
}
|
||||
|
||||
func prepareVAD(vad *config.VAD, modelsDir string) error {
|
||||
if vad == nil || !vad.Enabled {
|
||||
return nil
|
||||
}
|
||||
if modelsDir != "" {
|
||||
vad.Model = vad.ResolveModelPath(modelsDir)
|
||||
}
|
||||
*vad = vad.WithDefaults()
|
||||
if err := vad.Validate(); err != nil {
|
||||
return fmt.Errorf("vad: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go-whisper-api/config"
|
||||
|
||||
"github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
type OutputFormat string
|
||||
|
||||
func (f OutputFormat) String() string {
|
||||
return string(f)
|
||||
}
|
||||
|
||||
var (
|
||||
FormatTxt OutputFormat = "txt"
|
||||
FormatSrt OutputFormat = "srt"
|
||||
FormatCSV OutputFormat = "csv"
|
||||
)
|
||||
|
||||
type Engine struct {
|
||||
cfg *config.Whisper
|
||||
ctx whisper.Context
|
||||
model whisper.Model
|
||||
segments []whisper.Segment
|
||||
progress int
|
||||
runOpts RunOptions
|
||||
}
|
||||
|
||||
func (e *Engine) Transcript() error {
|
||||
return defaultPool.WithModel(e.cfg.Model, func(m whisper.Model) error {
|
||||
return e.transcribeWithModel(m)
|
||||
})
|
||||
}
|
||||
|
||||
func (e *Engine) transcribeWithModel(model whisper.Model) error {
|
||||
data, cleanup, err := prepareAudioPCM(e.cfg.AudioPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cleanup()
|
||||
e.model = model
|
||||
e.ctx, err = e.model.NewContext()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
e.ctx.SetThreads(e.cfg.Threads)
|
||||
if e.cfg.SpeedUp {
|
||||
e.ctx.SetAudioCtx(750)
|
||||
}
|
||||
e.ctx.SetTranslate(e.cfg.Translate)
|
||||
if e.cfg.Prompt != "" {
|
||||
e.ctx.SetInitialPrompt(e.cfg.Prompt)
|
||||
}
|
||||
e.ctx.SetMaxContext(int(e.cfg.MaxContext))
|
||||
if e.cfg.Debug {
|
||||
log.Info().Msgf("%s", e.ctx.SystemInfo())
|
||||
}
|
||||
if e.cfg.Language != "" {
|
||||
_ = e.ctx.SetLanguage(e.cfg.Language)
|
||||
}
|
||||
if e.cfg.BeamSize > 0 {
|
||||
e.ctx.SetBeamSize(int(e.cfg.BeamSize))
|
||||
}
|
||||
if e.cfg.EntropyThold > 0 {
|
||||
e.ctx.SetEntropyThold(float32(e.cfg.EntropyThold))
|
||||
}
|
||||
if err := prepareVAD(&e.cfg.VAD, ""); err != nil {
|
||||
return err
|
||||
}
|
||||
ApplyVAD(e.ctx, e.cfg.VAD)
|
||||
log.Debug().Msg("start transcribe process")
|
||||
e.ctx.ResetTimings()
|
||||
if err := e.ctx.Process(data, e.cbEncoderBegin(), e.cbSegment(), e.cbProgress()); err != nil {
|
||||
return err
|
||||
}
|
||||
if e.cfg.Debug {
|
||||
e.ctx.PrintTimings()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *Engine) cbEncoderBegin() func() bool {
|
||||
return func() bool { return true }
|
||||
}
|
||||
|
||||
func (e *Engine) cbSegment() func(segment whisper.Segment) {
|
||||
return func(segment whisper.Segment) {
|
||||
e.segments = append(e.segments, segment)
|
||||
if !e.cfg.PrintSegment {
|
||||
return
|
||||
}
|
||||
log.Info().Msgf(
|
||||
"[%6s -> %6s] %s",
|
||||
segment.Start.Truncate(time.Millisecond),
|
||||
segment.End.Truncate(time.Millisecond),
|
||||
segment.Text,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) cbProgress() func(progress int) {
|
||||
return func(progress int) {
|
||||
if progress > 100 {
|
||||
progress = 100
|
||||
}
|
||||
if e.progress == progress {
|
||||
return
|
||||
}
|
||||
e.progress = progress
|
||||
if e.cfg.PrintProgress {
|
||||
log.Info().Msgf("current progress: %d%%", progress)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) getOutputPath(format string) string {
|
||||
ext := filepath.Ext(e.cfg.AudioPath)
|
||||
filename := filepath.Base(e.cfg.AudioPath)
|
||||
if e.cfg.OutputFilename != "" {
|
||||
filename = e.cfg.OutputFilename
|
||||
}
|
||||
folder := filepath.Dir(e.cfg.AudioPath)
|
||||
if e.cfg.OutputFolder != "" {
|
||||
folder = e.cfg.OutputFolder
|
||||
}
|
||||
return path.Join(folder, strings.TrimSuffix(filename, ext)+"."+format)
|
||||
}
|
||||
|
||||
func (e *Engine) Save(format string) error {
|
||||
outputPath := e.getOutputPath(format)
|
||||
log.Info().Str("output-path", outputPath).Str("output-format", format).Msg("save text to file")
|
||||
text := ""
|
||||
switch OutputFormat(format) {
|
||||
case FormatSrt:
|
||||
for i, segment := range e.segments {
|
||||
text += fmt.Sprintf("%d\n", i+1)
|
||||
text += fmt.Sprintf("%s --> %s\n", srtTimestamp(segment.Start), srtTimestamp(segment.End))
|
||||
text += segment.Text + "\n\n"
|
||||
|
||||
}
|
||||
case FormatTxt:
|
||||
for _, segment := range e.segments {
|
||||
text += segment.Text
|
||||
}
|
||||
case FormatCSV:
|
||||
text = "start,end,text\n"
|
||||
for _, segment := range e.segments {
|
||||
text += fmt.Sprintf("%s,%s,\"%s\"\n", segment.Start, segment.End, segment.Text)
|
||||
}
|
||||
}
|
||||
if err := os.WriteFile(outputPath, []byte(text), 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type Word struct {
|
||||
Word string `json:"word"`
|
||||
Start int `json:"start"`
|
||||
Stop int `json:"stop"`
|
||||
}
|
||||
|
||||
type TranscriptResult struct {
|
||||
Text string `json:"text"`
|
||||
Words []Word `json:"words,omitempty"`
|
||||
}
|
||||
|
||||
func (e *Engine) SetTranscriptText(text string) {
|
||||
if len(e.segments) == 0 {
|
||||
e.segments = []whisper.Segment{{Text: text}}
|
||||
return
|
||||
}
|
||||
start := e.segments[0].Start
|
||||
end := e.segments[len(e.segments)-1].End
|
||||
e.segments = []whisper.Segment{{Text: text, Start: start, End: end}}
|
||||
}
|
||||
|
||||
func (e *Engine) Result() TranscriptResult {
|
||||
segments := e.segments
|
||||
if e.runOpts.PunctuateRestore != nil {
|
||||
updated, err := PunctuateSegments(segments, e.runOpts.PunctuateRestore)
|
||||
if err == nil {
|
||||
segments = updated
|
||||
}
|
||||
}
|
||||
text := FormatSegments(segments, e.runOpts.Turns, e.runOpts.Format)
|
||||
var words []Word
|
||||
for _, segment := range segments {
|
||||
words = append(words, segmentWords(segment)...)
|
||||
}
|
||||
return TranscriptResult{
|
||||
Text: text,
|
||||
Words: words,
|
||||
}
|
||||
}
|
||||
|
||||
func segmentWords(segment whisper.Segment) []Word {
|
||||
parts := strings.Fields(strings.TrimSpace(segment.Text))
|
||||
if len(parts) == 0 {
|
||||
return nil
|
||||
}
|
||||
startMs := int(segment.Start / time.Millisecond)
|
||||
endMs := int(segment.End / time.Millisecond)
|
||||
if endMs < startMs {
|
||||
endMs = startMs
|
||||
}
|
||||
span := endMs - startMs
|
||||
if span <= 0 {
|
||||
span = 1
|
||||
}
|
||||
step := span / len(parts)
|
||||
if step < 1 {
|
||||
step = 1
|
||||
}
|
||||
out := make([]Word, 0, len(parts))
|
||||
for i, part := range parts {
|
||||
wStart := startMs + i*step
|
||||
wStop := wStart + step
|
||||
if i == len(parts)-1 {
|
||||
wStop = endMs
|
||||
}
|
||||
out = append(out, Word{
|
||||
Word: part,
|
||||
Start: wStart,
|
||||
Stop: wStop,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *Engine) Close() error {
|
||||
// Models are owned by ModelPool; do not close shared weights here.
|
||||
e.ctx = nil
|
||||
e.model = nil
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package whisper
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"go-whisper-api/config"
|
||||
|
||||
"github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
|
||||
)
|
||||
|
||||
func TestEngine_getOutputPath(t *testing.T) {
|
||||
type fields struct {
|
||||
cfg *config.Whisper
|
||||
ctx whisper.Context
|
||||
model whisper.Model
|
||||
segments []whisper.Segment
|
||||
}
|
||||
type args struct {
|
||||
format string
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
fields fields
|
||||
args args
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "change wav to txt",
|
||||
fields: fields{
|
||||
cfg: &config.Whisper{
|
||||
AudioPath: "/test/1234/foo.wav",
|
||||
},
|
||||
},
|
||||
args: args{
|
||||
format: "txt",
|
||||
},
|
||||
want: "/test/1234/foo.txt",
|
||||
},
|
||||
{
|
||||
name: "change output folder",
|
||||
fields: fields{
|
||||
cfg: &config.Whisper{
|
||||
AudioPath: "/test/1234/foo.wav",
|
||||
OutputFolder: "/foo/bar",
|
||||
},
|
||||
},
|
||||
args: args{
|
||||
format: "txt",
|
||||
},
|
||||
want: "/foo/bar/foo.txt",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
e := &Engine{
|
||||
cfg: tt.fields.cfg,
|
||||
ctx: tt.fields.ctx,
|
||||
model: tt.fields.model,
|
||||
segments: tt.fields.segments,
|
||||
}
|
||||
if got := e.getOutputPath(tt.args.format); got != tt.want {
|
||||
t.Errorf("Engine.getOutputPath() = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user