Repository: ollama/ollama
File: integration/embed_test.go (Go)

Code before:
	return sum
}

func TestCosineDistanceCorrelation(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	client, _, cleanup := InitServerConnection(ctx, t)
	defer cleanup()

	for _, model := range libraryEmbedModels {
		testCases := []struct {
			a string
			b string
			c string
		}{
			{"cat", "kitten", "dog"},
			{"king", "queen", "baron"},
			{"paris", "london", "vancouver"},
			{"The cat is sleeping on the sofa", "A feline is sleeping on the couch", "Quantum physics is complex"},
			{"I love programming in python", "Coding in python brings me joy", "Pizza is delicious"},
			{"Machine learning is fascinating", "Artificial intelligence is amazing", "I need to buy groceries"},
			{"The quick brown fox jumps over the lazy dog", "A fast brown fox leaps over a sleepy dog", "The weather is warm and sunny today"},
		}

		for _, tc := range testCases {
			testEmbed := make(map[string][]float32)
			strs := []string{tc.a, tc.b, tc.c}

			req := api.EmbedRequest{
				Model:     model,
				Input:     strs,
				KeepAlive: &api.Duration{Duration: 10 * time.Second},
			}

			resp, err := embedTestHelper(ctx, client, t, req)
			if err != nil {
				t.Fatal(err)
			}

			for cnt, v := range resp.Embeddings {
				testEmbed[strs[cnt]] = v
			}

			// Calculate cosine similarities
			cosAB := cosineSimilarity(testEmbed[tc.a], testEmbed[tc.b])
			cosAC := cosineSimilarity(testEmbed[tc.a], testEmbed[tc.c])

			// Calculate distances
			distAB := euclideanDistance(testEmbed[tc.a], testEmbed[tc.b])
			distAC := euclideanDistance(testEmbed[tc.a], testEmbed[tc.c])


Review comment (suggestion):
It might be useful to inject a `t.Run(model, func(t *testing.T) {` in here so the discrete models could be tested with a `-run TestCosineDistanceCorrelation/embeddinggemma`

Code after:

func TestEmbedCosineDistanceCorrelation(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	client, _, cleanup := InitServerConnection(ctx, t)
	defer cleanup()

	for _, model := range libraryEmbedModels {
		t.Run(model, func(t *testing.T) {
			testCases := []struct {
				a string
				b string
				c string
			}{
				{"cat", "kitten", "dog"},
				{"king", "queen", "baron"},
				{"paris", "london", "vancouver"},
				{"The cat is sleeping on the sofa", "A feline is sleeping on the couch", "Quantum physics is complex"},
				{"I love programming in python", "Coding in python brings me joy", "Pizza is delicious"},
				{"Machine learning is fascinating", "Artificial intelligence is amazing", "I need to buy groceries"},
				{"The quick brown fox jumps over the lazy dog", "A fast brown fox leaps over a sleepy dog", "The weather is warm and sunny today"},
			}

			for _, tc := range testCases {
				testEmbed := make(map[string][]float32)
				strs := []string{tc.a, tc.b, tc.c}

				req := api.EmbedRequest{
					Model:     model,
					Input:     strs,
					KeepAlive: &api.Duration{Duration: 10 * time.Second},
				}

				resp, err := embedTestHelper(ctx, client, t, req)
				if err != nil {
					t.Fatal(err)
				}

				for cnt, v := range resp.Embeddings {
					testEmbed[strs[cnt]] = v
				}

				// Calculate cosine similarities
				cosAB := cosineSimilarity(testEmbed[tc.a], testEmbed[tc.b])
				cosAC := cosineSimilarity(testEmbed[tc.a], testEmbed[tc.c])

				// Calculate distances
				distAB := euclideanDistance(testEmbed[tc.a], testEmbed[tc.b])
				distAC := euclideanDistance(testEmbed[tc.a], testEmbed[tc.c])

				manhattanAB := manhattanDistance(testEmbed[tc.a], testEmbed[tc.b])