refactor: added segment formatiing, fixed benchmark printings
This commit is contained in:
parent
50caee3467
commit
91f5cc12f7
|
|
@ -2,6 +2,7 @@ package whisper_test
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"io"
|
||||||
"math"
|
"math"
|
||||||
"os"
|
"os"
|
||||||
"runtime"
|
"runtime"
|
||||||
|
|
@ -13,6 +14,38 @@ import (
|
||||||
wav "github.com/go-audio/wav"
|
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,
|
// benchProcessVariants runs the common benchmark matrix across context kinds,
|
||||||
// thread sets, and callback modes, for given samples. If singleIteration is true
|
// 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,
|
// it runs only one iteration regardless of b.N. If printTimings is true,
|
||||||
|
|
@ -87,18 +120,23 @@ func benchProcessVariants(
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < iters; i++ {
|
for i := 0; i < iters; i++ {
|
||||||
if printTimings {
|
model.ResetTimings()
|
||||||
model.ResetTimings()
|
|
||||||
}
|
|
||||||
start := time.Now()
|
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 {
|
if printTimings {
|
||||||
elapsed := time.Since(start)
|
|
||||||
model.PrintTimings()
|
model.PrintTimings()
|
||||||
b.ReportMetric(float64(elapsed.Milliseconds()), "ms_process")
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
b.ReportMetric(float64(elapsed.Milliseconds()), "ms_process")
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
|
|
||||||
|
|
@ -120,14 +158,22 @@ func benchProcessVariants(
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < iters; i++ {
|
for i := 0; i < iters; i++ {
|
||||||
start := time.Now()
|
start := time.Now()
|
||||||
|
model.ResetTimings()
|
||||||
|
|
||||||
// Passing a segment callback forces single-segment mode and exercises token extraction
|
// 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.Fatalf("process with callback: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
b.Logf("segments: %+v", segments)
|
||||||
|
|
||||||
|
elapsed := time.Since(start)
|
||||||
if printTimings {
|
if printTimings {
|
||||||
elapsed := time.Since(start)
|
model.PrintTimings()
|
||||||
b.ReportMetric(float64(elapsed.Milliseconds()), "ms_process")
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
b.ReportMetric(float64(elapsed.Milliseconds()), "ms_process")
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
package whisper
|
package whisper
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
@ -175,6 +176,11 @@ type Segment struct {
|
||||||
SpeakerTurnNext bool
|
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
|
// Token is a text or special token
|
||||||
type Token struct {
|
type Token struct {
|
||||||
// ID of the token
|
// ID of the token
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue