From 91f5cc12f77f925c9b8a0fe8f577a0e9012faf18 Mon Sep 17 00:00:00 2001
From: ciricc <facevk.02@mail.ru>
Date: Sat, 20 Sep 2025 17:03:25 +0300
Subject: [PATCH] refactor: added segment formatiing, fixed benchmark printings

---
 .../go/pkg/whisper/context_benchmark_test.go  | 66 ++++++++++++++++---
 bindings/go/pkg/whisper/interface.go          |  6 ++
 2 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/bindings/go/pkg/whisper/context_benchmark_test.go b/bindings/go/pkg/whisper/context_benchmark_test.go
index 04cf5c897..8cc6e5d30 100644
--- a/bindings/go/pkg/whisper/context_benchmark_test.go
+++ b/bindings/go/pkg/whisper/context_benchmark_test.go
@@ -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")
 					}
 				})
 			}
diff --git a/bindings/go/pkg/whisper/interface.go b/bindings/go/pkg/whisper/interface.go
index 4bd0262be..eabdb2db0 100644
--- a/bindings/go/pkg/whisper/interface.go
+++ b/bindings/go/pkg/whisper/interface.go
@@ -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