Merge pull request #3 from lein3000zzz/lein3000zzz

Fix tests

Author: clems4ever

all_minilm_l6_v2/benchmark_test.go

@@ -14,6 +14,7 @@ func newModel(t *testing.B) *all_minilm_l6_v2.Model {
 	}
 	t.Cleanup(func() {
 		benchModel.Close()
+		all_minilm_l6_v2.DestroyOrtEnv()
 	})
 	return benchModel
 }

all_minilm_l6_v2/model.go

@@ -102,6 +102,11 @@ func (m *Model) Close() {
 	}
 }
 
+func DestroyOrtEnv() error {
+	err := ort.DestroyEnvironment()
+	return err
+}
+
 func (m *Model) Compute(sentence string, addSpecialTokens bool) ([]float32, error) {
 	results, err := m.ComputeBatch([]string{sentence}, addSpecialTokens)
 	if err != nil {

all_minilm_l6_v2/model_test.go

@@ -12,7 +12,7 @@ func TestSingleSentenceEmbedding(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create model: %v", err)
 	}
-	defer model.Close()
+	defer closeAndDestroy(model)
 
 	sentence := "Hello, world! This is a test sentence."
 	embedding, err := model.Compute(sentence, false)
@@ -54,7 +54,7 @@ func TestBatchEmbedding(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create model: %v", err)
 	}
-	defer model.Close()
+	defer closeAndDestroy(model)
 
 	sentences := []string{
 		"Hello, world! This is a test sentence.",
@@ -117,7 +117,7 @@ func TestConsistentEmbeddingForSameSentence(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create model: %v", err)
 	}
-	defer model.Close()
+	defer closeAndDestroy(model)
 
 	// Test with same sentence appearing twice in batch, plus a different sentence
 	sentences := []string{
@@ -172,7 +172,7 @@ func TestSingleVsBatchConsistency(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create model: %v", err)
 	}
-	defer model.Close()
+	defer closeAndDestroy(model)
 
 	sentence := "Testing consistency between single and batch computation."
 
@@ -203,7 +203,7 @@ func TestEmptyBatch(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create model: %v", err)
 	}
-	defer model.Close()
+	defer closeAndDestroy(model)
 
 	// Test empty batch
 	embeddings, err := model.ComputeBatch([]string{}, false)
@@ -230,3 +230,9 @@ func vectorsEqual(a, b []float32) bool {
 	}
 	return true
 }
+
+// Helper function to close the model and destroy its env
+func closeAndDestroy(m *all_minilm_l6_v2.Model) {
+	m.Close()
+	all_minilm_l6_v2.DestroyOrtEnv()
+}