Update query_rewriter.go

corrected sqrt in euclidean and cosine
hypermodeinc · May 8, 2024 · 6544286 · 6544286
1 parent f1ece2b
commit 6544286
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/graphql/resolve/query_rewriter.go b/graphql/resolve/query_rewriter.go
@@ -659,12 +659,12 @@ func rewriteAsSimilarByIdQuery(
 	topK := query.ArgValue(schema.SimilarTopKArgName)
 	similarByField := typ.Field(similarBy)
 	metric := similarByField.EmbeddingSearchMetric()
-	distanceFormula := "math((v2 - v1) dot (v2 - v1))" // default - euclidian
+	distanceFormula := "math(sqrt((v2 - v1) dot (v2 - v1)))" // default - euclidian
 
 	if metric == schema.SimilarSearchMetricDotProduct {
 		distanceFormula = "math((1.0 - (v1 dot v2)) /2.0)"
 	} else if metric == schema.SimilarSearchMetricCosine {
-		distanceFormula = "math((1.0 - ((v1 dot v2) / ((v1 dot v1) * (v2 dot v2)))) / 2.0)"
+		distanceFormula = "math((1.0 - ((v1 dot v2) / sqrt( (v1 dot v1) * (v2 dot v2) ) )) / 2.0)"
 	}
 
 	// First generate the query to fetch the uid
@@ -819,12 +819,12 @@ func rewriteAsSimilarByEmbeddingQuery(
 
 	similarByField := typ.Field(similarBy)
 	metric := similarByField.EmbeddingSearchMetric()
-	distanceFormula := "math((v2 - $search_vector) dot (v2 - $search_vector))" // default = euclidian
+	distanceFormula := "math(sqrt((v2 - $search_vector) dot (v2 - $search_vector)))" // default = euclidian
 
 	if metric == schema.SimilarSearchMetricDotProduct {
 		distanceFormula = "math(( 1.0 - (($search_vector) dot v2)) /2.0)"
 	} else if metric == schema.SimilarSearchMetricCosine {
-		distanceFormula = "math(( ($search_vector) dot v2) / (( ($search_vector) dot ($search_vector)) * (v2 dot v2)))"
+		distanceFormula = "math((1.0 - ( (($search_vector) dot v2) / sqrt( (($search_vector) dot ($search_vector)) * (v2 dot v2) ) )) / 2.0)"
 	}
 
 	// Save vectorString as a query variable, $search_vector