Merge pull request #6 from knights-analytics/interface-test

Add RunPipeline function for concrete return types

Author: riccardopinosio

README.md

@@ -85,6 +85,11 @@ import (
 	"github.com/knights-analytics/hugot/pipelines"
 )
 
+func check(err error) {
+    if err != nil {
+        panic(err.Error())
+    }
+}
 // start a new session. This looks for the onnxruntime.so library in its default path, e.g. /usr/lib/onnxruntime.so
 session, err := hugot.NewSession()
 // if your onnxruntime.so is somewhere else, you can explicitly set it by using WithOnnxLibraryPath
@@ -97,17 +102,14 @@ defer func(session *hugot.Session) {
 }(session)
 // we now create a text classification pipeline. It requires the path to the onnx model folder,
 // and a pipeline name
-sentimentPipeline, err := session.NewTextClassificationPipeline(modelPath, "testPipeline")
+sentimentPipeline, err := session.NewTextClassificationPipeline("/path/to/model/", "testPipeline")
 check(err)
 // we can now use the pipeline for prediction on a batch of strings
 batch := []string{"This movie is disgustingly good !", "The director tried too much"}
-batchResult, err := sentimentPipeline.Run(batch)
+batchResult, err := sentimentPipeline.RunPipeline(batch)
 check(err)
-// batchResult is an interface so that we can treat pipelines uniformly.
-// we can cast it to the concrete result type of this pipeline
-result, ok := batchResult.(*pipelines.TextClassificationOutput)
 // and do whatever we want with it :)
-s, err := json.Marshal(result)
+s, err := json.Marshal(batchResult)
 check(err)
 fmt.Println(string(s))
 // {"ClassificationOutputs":[[{"Label":"POSITIVE","Score":0.9998536}],[{"Label":"NEGATIVE","Score":0.99752176}]]}
@@ -208,4 +210,4 @@ This will build a test image and run all tests in a container. A testTarget fold
 3. add tests and make sure the full test suite passes and test coverage does not dip below 80%
 4. create a MR linking to the relevant issue
 
-Thank you for contributing to hugot!
+Thank you for contributing to hugot!

hugot_test.go

@@ -73,14 +73,10 @@ func TestTextClassificationPipeline(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			batchResult, err := tt.pipeline.Run(tt.strings)
-			result, ok := batchResult.(*pipelines.TextClassificationOutput)
-			if !ok {
-				t.FailNow()
-			}
+			batchResult, err := tt.pipeline.RunPipeline(tt.strings)
 			check(t, err)
 			for i, expected := range tt.expected.ClassificationOutputs {
-				checkClassificationOutput(t, expected, result.ClassificationOutputs[i])
+				checkClassificationOutput(t, expected, batchResult.ClassificationOutputs[i])
 			}
 		})
 	}
@@ -176,15 +172,10 @@ func TestTokenClassificationPipeline(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			batchResult, err := tt.pipeline.Run(tt.strings)
+			batchResult, err := tt.pipeline.RunPipeline(tt.strings)
 			check(t, err)
-			result, ok := batchResult.(*pipelines.TokenClassificationOutput)
-			if !ok {
-				t.FailNow()
-			}
-
-			printTokenEntities(result)
-			for i, predictedEntities := range result.Entities {
+			printTokenEntities(batchResult)
+			for i, predictedEntities := range batchResult.Entities {
 				assert.Equal(t, len(tt.expected.Entities[i]), len(predictedEntities))
 				for j, entity := range predictedEntities {
 					expectedEntity := tt.expected.Entities[i][j]
@@ -254,13 +245,9 @@ func TestFeatureExtractionPipeline(t *testing.T) {
 	// test 'robert smith'
 	testResults = expectedResults["test1output"]
 	for i := 1; i <= 10; i++ {
-		batchResult, err := pipeline.Run([]string{"robert smith"})
+		batchResult, err := pipeline.RunPipeline([]string{"robert smith"})
 		check(t, err)
-		result, ok := batchResult.(*pipelines.FeatureExtractionOutput)
-		if !ok {
-			t.FailNow()
-		}
-		e := floatsEqual(result.Embeddings[0], testResults[0])
+		e := floatsEqual(batchResult.Embeddings[0], testResults[0])
 		if e != nil {
 			t.Logf("Test 1: The neural network didn't produce the correct result on loop %d: %s\n", i, e)
 			t.FailNow()
@@ -270,13 +257,9 @@ func TestFeatureExtractionPipeline(t *testing.T) {
 	// test ['robert smith junior', 'francis ford coppola']
 	testResults = expectedResults["test2output"]
 	for i := 1; i <= 10; i++ {
-		batchResult, err := pipeline.Run([]string{"robert smith junior", "francis ford coppola"})
+		batchResult, err := pipeline.RunPipeline([]string{"robert smith junior", "francis ford coppola"})
 		check(t, err)
-		result, ok := batchResult.(*pipelines.FeatureExtractionOutput)
-		if !ok {
-			t.FailNow()
-		}
-		for j, res := range result.Embeddings {
+		for j, res := range batchResult.Embeddings {
 			e := floatsEqual(res, testResults[j])
 			if e != nil {
 				t.Logf("Test 2: The neural network didn't produce the correct result on loop %d: %s\n", i, e)
@@ -293,21 +276,13 @@ func TestFeatureExtractionPipeline(t *testing.T) {
 
 	for k, sentencePair := range testPairs {
 		// these vectors should be the same
-		firstBatchResult, err2 := pipeline.Run(sentencePair[0])
+		firstBatchResult, err2 := pipeline.RunPipeline(sentencePair[0])
 		check(t, err2)
-		firstResult, ok := firstBatchResult.(*pipelines.FeatureExtractionOutput)
-		if !ok {
-			t.FailNow()
-		}
-		firstEmbedding := firstResult.Embeddings[0]
+		firstEmbedding := firstBatchResult.Embeddings[0]
 
-		secondBatchResult, err3 := pipeline.Run(sentencePair[1])
+		secondBatchResult, err3 := pipeline.RunPipeline(sentencePair[1])
 		check(t, err3)
-		secondResult, ok := secondBatchResult.(*pipelines.FeatureExtractionOutput)
-		if !ok {
-			t.FailNow()
-		}
-		secondEmbedding := secondResult.Embeddings[0]
+		secondEmbedding := secondBatchResult.Embeddings[0]
 		e := floatsEqual(firstEmbedding, secondEmbedding)
 		if e != nil {
 			t.Logf("Equality failed for determinism test %s test with pairs %s and %s", k, strings.Join(sentencePair[0], ","), strings.Join(sentencePair[1], ","))

pipelines/featureExtraction.go

@@ -66,7 +66,7 @@ func (p *FeatureExtractionPipeline) Validate() error {
 }
 
 // Postprocess Parse the results of the forward pass into the output. Token embeddings are mean pooled.
-func (p *FeatureExtractionPipeline) Postprocess(batch PipelineBatch) (PipelineBatchOutput, error) {
+func (p *FeatureExtractionPipeline) Postprocess(batch PipelineBatch) (*FeatureExtractionOutput, error) {
 
 	maxSequence := batch.MaxSequence
 	vectorCounter := 0
@@ -119,6 +119,10 @@ func meanPooling(tokens [][]float32, input TokenizedInput, maxSequence int, dime
 
 // Run the pipeline on a string batch
 func (p *FeatureExtractionPipeline) Run(inputs []string) (PipelineBatchOutput, error) {
+	return p.RunPipeline(inputs)
+}
+
+func (p *FeatureExtractionPipeline) RunPipeline(inputs []string) (*FeatureExtractionOutput, error) {
 	batch := p.Preprocess(inputs)
 	batch, forwardError := p.Forward(batch)
 	if forwardError != nil {

pipelines/textClassification.go

@@ -159,7 +159,7 @@ func (p *TextClassificationPipeline) Forward(batch PipelineBatch) (PipelineBatch
 	return batch, err
 }
 
-func (p *TextClassificationPipeline) Postprocess(batch PipelineBatch) (PipelineBatchOutput, error) {
+func (p *TextClassificationPipeline) Postprocess(batch PipelineBatch) (*TextClassificationOutput, error) {
 
 	outputTensor := batch.OutputTensor
 	output := make([][]float32, len(batch.Input))
@@ -210,6 +210,10 @@ func (p *TextClassificationPipeline) Postprocess(batch PipelineBatch) (PipelineB
 
 // Run the pipeline on a string batch
 func (p *TextClassificationPipeline) Run(inputs []string) (PipelineBatchOutput, error) {
+	return p.RunPipeline(inputs)
+}
+
+func (p *TextClassificationPipeline) RunPipeline(inputs []string) (*TextClassificationOutput, error) {
 	batch := p.Preprocess(inputs)
 	batch, err := p.Forward(batch)
 	if err != nil {

pipelines/tokenClassification.go

@@ -149,7 +149,7 @@ func (p *TokenClassificationPipeline) Validate() error {
 }
 
 // Postprocess function for a token classification pipeline
-func (p *TokenClassificationPipeline) Postprocess(batch PipelineBatch) (PipelineBatchOutput, error) {
+func (p *TokenClassificationPipeline) Postprocess(batch PipelineBatch) (*TokenClassificationOutput, error) {
 
 	outputs := make([][][]float32, len(batch.Input))        // holds the final output
 	inputVectors := make([][]float32, 0, batch.MaxSequence) // holds the embeddings of each original token (no padding) for an input
@@ -355,6 +355,10 @@ func (p *TokenClassificationPipeline) GroupEntities(entities []Entity) ([]Entity
 
 // Run the pipeline on a string batch
 func (p *TokenClassificationPipeline) Run(inputs []string) (PipelineBatchOutput, error) {
+	return p.RunPipeline(inputs)
+}
+
+func (p *TokenClassificationPipeline) RunPipeline(inputs []string) (*TokenClassificationOutput, error) {
 	batch := p.Preprocess(inputs)
 	batch, errForward := p.Forward(batch)
 	if errForward != nil {