refactor: added segment formatiing, fixed benchmark printings

This commit is contained in:
ciricc 2025-09-20 17:03:25 +03:00
parent 50caee3467
commit 91f5cc12f7
2 changed files with 62 additions and 10 deletions

View File

@ -2,6 +2,7 @@ package whisper_test
import (
"fmt"
"io"
"math"
"os"
"runtime"
@ -13,6 +14,38 @@ import (
wav "github.com/go-audio/wav"
)
func processAndExtractSegmentsSequentially(ctx whisper.Context, samples []float32) ([]whisper.Segment, error) {
if err := ctx.Process(samples, nil, nil, nil); err != nil {
return nil, err
}
var segments []whisper.Segment
for {
seg, err := ctx.NextSegment()
if err == io.EOF {
break
} else if err != nil {
return nil, err
}
segments = append(segments, seg)
}
return segments, nil
}
func processAndExtractSegmentsWithCallback(ctx whisper.Context, samples []float32) ([]whisper.Segment, error) {
segments := make([]whisper.Segment, 0)
if err := ctx.Process(samples, nil, func(seg whisper.Segment) {
segments = append(segments, seg)
}, nil); err != nil {
return nil, err
}
return segments, nil
}
// benchProcessVariants runs the common benchmark matrix across context kinds,
// thread sets, and callback modes, for given samples. If singleIteration is true
// it runs only one iteration regardless of b.N. If printTimings is true,
@ -87,18 +120,23 @@ func benchProcessVariants(
b.ResetTimer()
for i := 0; i < iters; i++ {
if printTimings {
model.ResetTimings()
}
model.ResetTimings()
start := time.Now()
if err := ctx.Process(samples, nil, nil, nil); err != nil {
b.Fatalf("process: %v", err)
segments, err := processAndExtractSegmentsSequentially(ctx, samples)
if err != nil {
b.Fatalf("process and extract segments sequentially: %v", err)
}
b.Logf("segments: %+v", segments)
elapsed := time.Since(start)
if printTimings {
elapsed := time.Since(start)
model.PrintTimings()
b.ReportMetric(float64(elapsed.Milliseconds()), "ms_process")
}
b.ReportMetric(float64(elapsed.Milliseconds()), "ms_process")
}
})
@ -120,14 +158,22 @@ func benchProcessVariants(
b.ResetTimer()
for i := 0; i < iters; i++ {
start := time.Now()
model.ResetTimings()
// Passing a segment callback forces single-segment mode and exercises token extraction
if err := ctx.Process(samples, nil, func(seg whisper.Segment) {}, nil); err != nil {
segments, err := processAndExtractSegmentsWithCallback(ctx, samples)
if err != nil {
b.Fatalf("process with callback: %v", err)
}
b.Logf("segments: %+v", segments)
elapsed := time.Since(start)
if printTimings {
elapsed := time.Since(start)
b.ReportMetric(float64(elapsed.Milliseconds()), "ms_process")
model.PrintTimings()
}
b.ReportMetric(float64(elapsed.Milliseconds()), "ms_process")
}
})
}

View File

@ -1,6 +1,7 @@
package whisper
import (
"fmt"
"io"
"time"
)
@ -175,6 +176,11 @@ type Segment struct {
SpeakerTurnNext bool
}
func (s Segment) String() string {
// foramt: [00:01:39.000 --> 00:01:50.000] And so, my fellow Americans, ask not what your country can do for you, ask what you can do for your country.
return fmt.Sprintf("[%s --> %s] %s", s.Start.Truncate(time.Millisecond), s.End.Truncate(time.Millisecond), s.Text)
}
// Token is a text or special token
type Token struct {
// ID of the token