make statistics a proper struct

riccardopinosio · RJKeevil · commit 593c620b8db4 · 2025-12-08T10:24:41.000+01:00
diff --git a/backends/pipeline.go b/backends/pipeline.go
@@ -1,10 +1,14 @@
 package backends
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
+	"time"
 
 	"github.com/knights-analytics/hugot/options"
+	"github.com/knights-analytics/hugot/util/safeconv"
 )
 
 // BasePipeline can be embedded by a pipeline.
@@ -53,13 +57,49 @@ type PipelineBatchOutput interface {
 
 // Pipeline is the interface that any pipeline must implement.
 type Pipeline interface {
-	GetStats() []string                        // Get the pipeline running stats
+	GetStatistics() PipelineStatistics         // Get the pipeline running stats
 	Validate() error                           // Validate the pipeline for correctness
 	GetMetadata() PipelineMetadata             // Return metadata information for the pipeline
 	GetModel() *Model                          // Return the model used by the pipeline
 	Run([]string) (PipelineBatchOutput, error) // Run the pipeline on an input
 }
 
+type PipelineStatistics struct {
+	TokenizerTotalTime      time.Duration
+	TokenizerExecutionCount uint64
+	TokenizerAvgQueryTime   time.Duration
+	OnnxTotalTime           time.Duration
+	OnnxExecutionCount      uint64
+	OnnxAvgQueryTime        time.Duration
+	TotalQueries            uint64
+	TotalDocuments          uint64
+	AverageLatency          time.Duration
+	AverageBatchSize        float64
+	FilteredResults         uint64
+}
+
+func (p *PipelineStatistics) ComputeTokenizerStatistics(timings *timings) {
+	p.TokenizerTotalTime = safeconv.U64ToDuration(timings.TotalNS)
+	p.TokenizerExecutionCount = timings.NumCalls
+	p.TokenizerAvgQueryTime = time.Duration(float64(timings.TotalNS) /
+		math.Max(1, float64(timings.NumCalls)))
+}
+
+func (p *PipelineStatistics) ComputeOnnxStatistics(timings *timings) {
+	p.OnnxTotalTime = safeconv.U64ToDuration(timings.TotalNS)
+	p.OnnxExecutionCount = timings.NumCalls
+	p.OnnxAvgQueryTime = time.Duration(float64(timings.TotalNS) /
+		math.Max(1, float64(timings.NumCalls)))
+}
+
+func (p *PipelineStatistics) Print() {
+	jsonData, err := json.MarshalIndent(p, "", "  ")
+	if err != nil {
+		fmt.Println(err)
+	}
+	fmt.Println(string(jsonData))
+}
+
 // PipelineOption is an option for a pipeline type.
 type PipelineOption[T Pipeline] func(eo T) error
 
diff --git a/hugot.go b/hugot.go
@@ -58,10 +58,10 @@ func newSession(backend string, opts ...options.WithOption) (*Session, error) {
 
 type pipelineMap[T backends.Pipeline] map[string]T
 
-func (m pipelineMap[T]) GetStats() []string {
-	var stats []string
+func (m pipelineMap[T]) GetStatistics() []backends.PipelineStatistics {
+	var stats []backends.PipelineStatistics
 	for _, p := range m {
-		stats = append(stats, p.GetStats()...)
+		stats = append(stats, p.GetStatistics())
 	}
 	return stats
 }
@@ -379,24 +379,33 @@ func (e *pipelineNotFoundError) Error() string {
 	return fmt.Sprintf("Pipeline with name %s not found", e.pipelineName)
 }
 
-// GetStats returns runtime statistics for all initialized pipelines for profiling purposes. We currently record for each pipeline:
+// GetStatistics returns runtime statistics for all initialized pipelines for profiling purposes. We currently record for each pipeline:
 // the total runtime of the tokenization step
 // the number of batch calls to the tokenization step
 // the average time per tokenization batch call
 // the total runtime of the inference (i.e. onnxruntime) step
 // the number of batch calls to the onnxruntime inference
 // the average time per onnxruntime inference batch call.
-func (s *Session) GetStats() []string {
+func (s *Session) GetStatistics() []backends.PipelineStatistics {
 	return slices.Concat(
-		s.tokenClassificationPipelines.GetStats(),
-		s.textClassificationPipelines.GetStats(),
-		s.featureExtractionPipelines.GetStats(),
-		s.zeroShotClassificationPipelines.GetStats(),
-		s.crossEncoderPipelines.GetStats(),
-		s.textGenerationPipelines.GetStats(),
+		s.tokenClassificationPipelines.GetStatistics(),
+		s.textClassificationPipelines.GetStatistics(),
+		s.featureExtractionPipelines.GetStatistics(),
+		s.imageClassificationPipelines.GetStatistics(),
+		s.zeroShotClassificationPipelines.GetStatistics(),
+		s.crossEncoderPipelines.GetStatistics(),
+		s.textGenerationPipelines.GetStatistics(),
 	)
 }
 
+// Print prints runtime statistics for all initialized pipelines to stdout.
+func (s *Session) PrintStatistics() {
+	stats := s.GetStatistics()
+	for _, v := range stats {
+		v.Print()
+	}
+}
+
 // Destroy deletes the hugot session and onnxruntime environment and all initialized pipelines, freeing memory.
 // A hugot session should be destroyed when not neeeded any more, preferably with a defer() call.
 func (s *Session) Destroy() error {
@@ -408,6 +417,7 @@ func (s *Session) Destroy() error {
 	s.featureExtractionPipelines = nil
 	s.tokenClassificationPipelines = nil
 	s.textClassificationPipelines = nil
+	s.imageClassificationPipelines = nil
 	s.zeroShotClassificationPipelines = nil
 	s.textGenerationPipelines = nil
 	s.crossEncoderPipelines = nil
diff --git a/hugot_test.go b/hugot_test.go
@@ -236,8 +236,8 @@ func textClassificationPipeline(t *testing.T, session *Session) {
 		}
 	})
 
-	// check get stats
-	session.GetStats()
+	// check print stats
+	session.PrintStatistics()
 }
 
 func textClassificationPipelineMulti(t *testing.T, session *Session) {
@@ -396,7 +396,10 @@ func textClassificationPipelineMulti(t *testing.T, session *Session) {
 	})
 
 	// check get stats
-	session.GetStats()
+	stats := session.GetStatistics()
+	for _, v := range stats {
+		v.Print()
+	}
 }
 
 func textClassificationPipelineValidation(t *testing.T, session *Session) {
diff --git a/pipelines/crossEncoder.go b/pipelines/crossEncoder.go
@@ -3,7 +3,6 @@ package pipelines
 import (
 	"errors"
 	"fmt"
-	"math"
 	"sort"
 	"sync/atomic"
 	"time"
@@ -104,27 +103,20 @@ func (p *CrossEncoderPipeline) GetMetadata() backends.PipelineMetadata {
 	}
 }
 
-func (p *CrossEncoderPipeline) GetStats() []string {
+func (p *CrossEncoderPipeline) GetStatistics() backends.PipelineStatistics {
 	avgLatency := p.stats.AverageLatency
 	if p.stats.TotalQueries > 0 {
 		avgLatency = time.Duration(float64(p.stats.AverageLatency) / float64(p.stats.TotalQueries))
 	}
-	return []string{
-		fmt.Sprintf("Statistics for pipeline: %s", p.PipelineName),
-		fmt.Sprintf("Total queries processed: %d", p.stats.TotalQueries),
-		fmt.Sprintf("Total documents scored: %d", p.stats.TotalDocuments),
-		fmt.Sprintf("Average latency per query: %s", avgLatency),
-		fmt.Sprintf("Average batch size: %.2f", p.stats.AverageBatchSize),
-		fmt.Sprintf("Filtered results: %d", p.stats.FilteredResults),
-		fmt.Sprintf("Tokenizer: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.Model.Tokenizer.TokenizerTimings.TotalNS),
-			p.Model.Tokenizer.TokenizerTimings.NumCalls,
-			time.Duration(float64(p.Model.Tokenizer.TokenizerTimings.TotalNS)/math.Max(1, float64(p.Model.Tokenizer.TokenizerTimings.NumCalls)))),
-		fmt.Sprintf("ONNX: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.PipelineTimings.TotalNS),
-			p.PipelineTimings.NumCalls,
-			time.Duration(float64(p.PipelineTimings.TotalNS)/math.Max(1, float64(p.PipelineTimings.NumCalls)))),
-	}
+	statistics := backends.PipelineStatistics{}
+	statistics.ComputeTokenizerStatistics(p.Model.Tokenizer.TokenizerTimings)
+	statistics.ComputeOnnxStatistics(p.PipelineTimings)
+	statistics.TotalQueries = p.stats.TotalQueries
+	statistics.TotalDocuments = p.stats.TotalDocuments
+	statistics.AverageLatency = avgLatency
+	statistics.AverageBatchSize = p.stats.AverageBatchSize
+	statistics.FilteredResults = p.stats.FilteredResults
+	return statistics
 }
 
 func (p *CrossEncoderPipeline) Validate() error {
diff --git a/pipelines/featureExtraction.go b/pipelines/featureExtraction.go
@@ -3,7 +3,6 @@ package pipelines
 import (
 	"errors"
 	"fmt"
-	"math"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -112,19 +111,12 @@ func (p *FeatureExtractionPipeline) GetMetadata() backends.PipelineMetadata {
 	}
 }
 
-// GetStats returns the runtime statistics for the pipeline.
-func (p *FeatureExtractionPipeline) GetStats() []string {
-	return []string{
-		fmt.Sprintf("Statistics for pipeline: %s", p.PipelineName),
-		fmt.Sprintf("Tokenizer: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.Model.Tokenizer.TokenizerTimings.TotalNS),
-			p.Model.Tokenizer.TokenizerTimings.NumCalls,
-			time.Duration(float64(p.Model.Tokenizer.TokenizerTimings.TotalNS)/math.Max(1, float64(p.Model.Tokenizer.TokenizerTimings.NumCalls)))),
-		fmt.Sprintf("ONNX: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.PipelineTimings.TotalNS),
-			p.PipelineTimings.NumCalls,
-			time.Duration(float64(p.PipelineTimings.TotalNS)/math.Max(1, float64(p.PipelineTimings.NumCalls)))),
-	}
+// GetStatistics returns the runtime statistics for the pipeline.
+func (p *FeatureExtractionPipeline) GetStatistics() backends.PipelineStatistics {
+	statistics := backends.PipelineStatistics{}
+	statistics.ComputeTokenizerStatistics(p.Model.Tokenizer.TokenizerTimings)
+	statistics.ComputeOnnxStatistics(p.PipelineTimings)
+	return statistics
 }
 
 // Validate checks that the pipeline is valid.
diff --git a/pipelines/imageClassification.go b/pipelines/imageClassification.go
@@ -6,7 +6,6 @@ import (
 	"image"
 	_ "image/jpeg" // adds jpeg support to image
 	_ "image/png"  // adds png support to image
-	"math"
 	"sort"
 	"strings"
 	"sync/atomic"
@@ -147,14 +146,10 @@ func (p *ImageClassificationPipeline) GetMetadata() backends.PipelineMetadata {
 	}
 }
 
-func (p *ImageClassificationPipeline) GetStats() []string {
-	return []string{
-		fmt.Sprintf("Statistics for pipeline: %s", p.PipelineName),
-		fmt.Sprintf("ONNX: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.PipelineTimings.TotalNS),
-			p.PipelineTimings.NumCalls,
-			time.Duration(float64(p.PipelineTimings.TotalNS)/math.Max(1, float64(p.PipelineTimings.NumCalls)))),
-	}
+func (p *ImageClassificationPipeline) GetStatistics() backends.PipelineStatistics {
+	statistics := backends.PipelineStatistics{}
+	statistics.ComputeOnnxStatistics(p.PipelineTimings)
+	return statistics
 }
 
 func (p *ImageClassificationPipeline) Validate() error {
diff --git a/pipelines/textClassification.go b/pipelines/textClassification.go
@@ -3,7 +3,6 @@ package pipelines
 import (
 	"errors"
 	"fmt"
-	"math"
 	"sync/atomic"
 	"time"
 
@@ -129,19 +128,12 @@ func (p *TextClassificationPipeline) GetMetadata() backends.PipelineMetadata {
 	}
 }
 
-// GetStats returns the runtime statistics for the pipeline.
-func (p *TextClassificationPipeline) GetStats() []string {
-	return []string{
-		fmt.Sprintf("Statistics for pipeline: %s", p.PipelineName),
-		fmt.Sprintf("Tokenizer: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.Model.Tokenizer.TokenizerTimings.TotalNS),
-			p.Model.Tokenizer.TokenizerTimings.NumCalls,
-			time.Duration(float64(p.Model.Tokenizer.TokenizerTimings.TotalNS)/math.Max(1, float64(p.Model.Tokenizer.TokenizerTimings.NumCalls)))),
-		fmt.Sprintf("ONNX: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.PipelineTimings.TotalNS),
-			p.PipelineTimings.NumCalls,
-			time.Duration(float64(p.PipelineTimings.TotalNS)/math.Max(1, float64(p.PipelineTimings.NumCalls)))),
-	}
+// GetStatistics returns the runtime statistics for the pipeline.
+func (p *TextClassificationPipeline) GetStatistics() backends.PipelineStatistics {
+	statistics := backends.PipelineStatistics{}
+	statistics.ComputeTokenizerStatistics(p.Model.Tokenizer.TokenizerTimings)
+	statistics.ComputeOnnxStatistics(p.PipelineTimings)
+	return statistics
 }
 
 // Validate checks that the pipeline is valid.
diff --git a/pipelines/textGeneration.go b/pipelines/textGeneration.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
-	"math"
 	"sync/atomic"
 	"text/template"
 	"time"
@@ -167,18 +166,12 @@ func (p *TextGenerationPipeline) GetModel() *backends.Model {
 	return p.Model
 }
 
-func (p *TextGenerationPipeline) GetStats() []string {
-	return []string{
-		fmt.Sprintf("Statistics for pipeline: %s", p.PipelineName),
-		fmt.Sprintf("Tokenizer: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.Model.Tokenizer.TokenizerTimings.TotalNS),
-			p.Model.Tokenizer.TokenizerTimings.NumCalls,
-			time.Duration(float64(p.Model.Tokenizer.TokenizerTimings.TotalNS)/math.Max(1, float64(p.Model.Tokenizer.TokenizerTimings.NumCalls)))),
-		fmt.Sprintf("ONNX: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.PipelineTimings.TotalNS),
-			p.PipelineTimings.NumCalls,
-			time.Duration(float64(p.PipelineTimings.TotalNS)/math.Max(1, float64(p.PipelineTimings.NumCalls)))),
-	}
+// GetStatistics returns the runtime statistics for the pipeline.
+func (p *TextGenerationPipeline) GetStatistics() backends.PipelineStatistics {
+	statistics := backends.PipelineStatistics{}
+	statistics.ComputeTokenizerStatistics(p.Model.Tokenizer.TokenizerTimings)
+	statistics.ComputeOnnxStatistics(p.PipelineTimings)
+	return statistics
 }
 
 func (p *TextGenerationPipeline) Validate() error {
diff --git a/pipelines/tokenClassification.go b/pipelines/tokenClassification.go
@@ -3,7 +3,6 @@ package pipelines
 import (
 	"errors"
 	"fmt"
-	"math"
 	"slices"
 	"strings"
 	"sync/atomic"
@@ -149,19 +148,12 @@ func (p *TokenClassificationPipeline) GetMetadata() backends.PipelineMetadata {
 	}
 }
 
-// GetStats returns the runtime statistics for the pipeline.
-func (p *TokenClassificationPipeline) GetStats() []string {
-	return []string{
-		fmt.Sprintf("Statistics for pipeline: %s", p.PipelineName),
-		fmt.Sprintf("Tokenizer: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.Model.Tokenizer.TokenizerTimings.TotalNS),
-			p.Model.Tokenizer.TokenizerTimings.NumCalls,
-			time.Duration(float64(p.Model.Tokenizer.TokenizerTimings.TotalNS)/math.Max(1, float64(p.Model.Tokenizer.TokenizerTimings.NumCalls)))),
-		fmt.Sprintf("ONNX: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.PipelineTimings.TotalNS),
-			p.PipelineTimings.NumCalls,
-			time.Duration(float64(p.PipelineTimings.TotalNS)/math.Max(1, float64(p.PipelineTimings.NumCalls)))),
-	}
+// GetStatistics returns the runtime statistics for the pipeline.
+func (p *TokenClassificationPipeline) GetStatistics() backends.PipelineStatistics {
+	statistics := backends.PipelineStatistics{}
+	statistics.ComputeTokenizerStatistics(p.Model.Tokenizer.TokenizerTimings)
+	statistics.ComputeOnnxStatistics(p.PipelineTimings)
+	return statistics
 }
 
 // Validate checks that the pipeline is valid.
diff --git a/pipelines/zeroShotClassification.go b/pipelines/zeroShotClassification.go
@@ -311,18 +311,12 @@ func (p *ZeroShotClassificationPipeline) GetMetadata() backends.PipelineMetadata
 	}
 }
 
-func (p *ZeroShotClassificationPipeline) GetStats() []string {
-	return []string{
-		fmt.Sprintf("Statistics for pipeline: %s", p.PipelineName),
-		fmt.Sprintf("Tokenizer: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.Model.Tokenizer.TokenizerTimings.TotalNS),
-			p.Model.Tokenizer.TokenizerTimings.NumCalls,
-			time.Duration(float64(p.Model.Tokenizer.TokenizerTimings.TotalNS)/math.Max(1, float64(p.Model.Tokenizer.TokenizerTimings.NumCalls)))),
-		fmt.Sprintf("ONNX: Total time=%s, Execution count=%d, Average query time=%s",
-			safeconv.U64ToDuration(p.PipelineTimings.TotalNS),
-			p.PipelineTimings.NumCalls,
-			time.Duration(float64(p.PipelineTimings.TotalNS)/math.Max(1, float64(p.PipelineTimings.NumCalls)))),
-	}
+// GetStatistics returns the runtime statistics for the pipeline.
+func (p *ZeroShotClassificationPipeline) GetStatistics() backends.PipelineStatistics {
+	statistics := backends.PipelineStatistics{}
+	statistics.ComputeTokenizerStatistics(p.Model.Tokenizer.TokenizerTimings)
+	statistics.ComputeOnnxStatistics(p.PipelineTimings)
+	return statistics
 }
 
 func (p *ZeroShotClassificationPipeline) Run(inputs []string) (backends.PipelineBatchOutput, error) {

Original file line number	Diff line number	Diff line change
`@@ -236,8 +236,8 @@ func textClassificationPipeline(t testing.T, session Session) {`
`236`	`236`	`}`
`237`	`237`	`})`
`238`	`238`
`239`		`- // check get stats`
`240`		`- session.GetStats()`
	`239`	`+ // check print stats`
	`240`	`+ session.PrintStatistics()`
`241`	`241`	`}`
`242`	`242`
`243`	`243`	`func textClassificationPipelineMulti(t testing.T, session Session) {`
`@@ -396,7 +396,10 @@ func textClassificationPipelineMulti(t testing.T, session Session) {`
`396`	`396`	`})`
`397`	`397`
`398`	`398`	`// check get stats`
`399`		`- session.GetStats()`
	`399`	`+ stats := session.GetStatistics()`
	`400`	`+ for _, v := range stats {`
	`401`	`+ v.Print()`
	`402`	`+ }`
`400`	`403`	`}`
`401`	`404`
`402`	`405`	`func textClassificationPipelineValidation(t testing.T, session Session) {`