refactor: added segment formatiing, fixed benchmark printings
This commit is contained in:
parent
50caee3467
commit
91f5cc12f7
|
|
@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue