summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorArne Juul <arnej@yahooinc.com>2023-06-26 12:29:45 +0000
committerArne Juul <arnej@yahooinc.com>2023-06-26 13:06:51 +0000
commitcc517d86dc886058cdc5f95a318945a6a328da28 (patch)
treecef1ca84628800fc8226ccb30626688a3efdf49b
parent626bcc6c265229d8c97f4e0a1c996013650b335e (diff)
add cosine_similarity
-rw-r--r--searchlib/abi-spec.json2
-rwxr-xr-xsearchlib/src/main/javacc/RankingExpressionParser.jj15
-rw-r--r--vespajlib/abi-spec.json16
-rw-r--r--vespajlib/src/main/java/com/yahoo/tensor/functions/CosineSimilarity.java63
-rw-r--r--vespajlib/src/test/java/com/yahoo/tensor/functions/CosineSimilarityTestCase.java66
5 files changed, 162 insertions, 0 deletions
diff --git a/searchlib/abi-spec.json b/searchlib/abi-spec.json
index 30f2cb5c6ea..7d6f2f8790c 100644
--- a/searchlib/abi-spec.json
+++ b/searchlib/abi-spec.json
@@ -947,6 +947,7 @@
"public final com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode tensorL1Normalize()",
"public final com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode tensorL2Normalize()",
"public final com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode tensorEuclideanDistance()",
+ "public final com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode tensorCosineSimilarity()",
"public final com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode tensorMatmul()",
"public final com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode tensorSoftmax()",
"public final com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode tensorXwPlusB()",
@@ -1100,6 +1101,7 @@
"public static final int L1_NORMALIZE",
"public static final int L2_NORMALIZE",
"public static final int EUCLIDEAN_DISTANCE",
+ "public static final int COSINE_SIMILARITY",
"public static final int MATMUL",
"public static final int SOFTMAX",
"public static final int XW_PLUS_B",
diff --git a/searchlib/src/main/javacc/RankingExpressionParser.jj b/searchlib/src/main/javacc/RankingExpressionParser.jj
index 744e629893e..41647a5ef5b 100755
--- a/searchlib/src/main/javacc/RankingExpressionParser.jj
+++ b/searchlib/src/main/javacc/RankingExpressionParser.jj
@@ -139,6 +139,7 @@ TOKEN :
<L1_NORMALIZE: "l1_normalize"> |
<L2_NORMALIZE: "l2_normalize"> |
<EUCLIDEAN_DISTANCE: "euclidean_distance"> |
+ <COSINE_SIMILARITY: "cosine_similarity"> |
<MATMUL: "matmul"> |
<SOFTMAX: "softmax"> |
<XW_PLUS_B: "xw_plus_b"> |
@@ -381,6 +382,7 @@ TensorFunctionNode tensorFunction() :
tensorExpression = tensorL1Normalize() |
tensorExpression = tensorL2Normalize() |
tensorExpression = tensorEuclideanDistance() |
+ tensorExpression = tensorCosineSimilarity() |
tensorExpression = tensorMatmul() |
tensorExpression = tensorSoftmax() |
tensorExpression = tensorXwPlusB() |
@@ -558,6 +560,18 @@ TensorFunctionNode tensorEuclideanDistance() :
dimension)); }
}
+TensorFunctionNode tensorCosineSimilarity() :
+{
+ ExpressionNode tensor1, tensor2;
+ String dimension;
+}
+{
+ <COSINE_SIMILARITY> <LBRACE> tensor1 = expression() <COMMA> tensor2 = expression() <COMMA> dimension = identifier() <RBRACE>
+ { return new TensorFunctionNode(new CosineSimilarity(TensorFunctionNode.wrap(tensor1),
+ TensorFunctionNode.wrap(tensor2),
+ dimension)); }
+}
+
TensorFunctionNode tensorMatmul() :
{
ExpressionNode tensor1, tensor2;
@@ -716,6 +730,7 @@ String tensorFunctionName() :
( <L1_NORMALIZE> { return token.image; } ) |
( <L2_NORMALIZE> { return token.image; } ) |
( <EUCLIDEAN_DISTANCE> { return token.image; } ) |
+ ( <COSINE_SIMILARITY> { return token.image; } ) |
( <MATMUL> { return token.image; } ) |
( <SOFTMAX> { return token.image; } ) |
( <XW_PLUS_B> { return token.image; } ) |
diff --git a/vespajlib/abi-spec.json b/vespajlib/abi-spec.json
index 3b9f494dc50..7f70deb0991 100644
--- a/vespajlib/abi-spec.json
+++ b/vespajlib/abi-spec.json
@@ -1705,6 +1705,22 @@
],
"fields" : [ ]
},
+ "com.yahoo.tensor.functions.CosineSimilarity" : {
+ "superClass" : "com.yahoo.tensor.functions.CompositeTensorFunction",
+ "interfaces" : [ ],
+ "attributes" : [
+ "public"
+ ],
+ "methods" : [
+ "public void <init>(com.yahoo.tensor.functions.TensorFunction, com.yahoo.tensor.functions.TensorFunction, java.lang.String)",
+ "public java.util.List arguments()",
+ "public com.yahoo.tensor.functions.TensorFunction withArguments(java.util.List)",
+ "public com.yahoo.tensor.functions.PrimitiveTensorFunction toPrimitive()",
+ "public java.lang.String toString(com.yahoo.tensor.functions.ToStringContext)",
+ "public int hashCode()"
+ ],
+ "fields" : [ ]
+ },
"com.yahoo.tensor.functions.Diag" : {
"superClass" : "com.yahoo.tensor.functions.CompositeTensorFunction",
"interfaces" : [ ],
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/CosineSimilarity.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/CosineSimilarity.java
new file mode 100644
index 00000000000..ede0355a3a6
--- /dev/null
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/CosineSimilarity.java
@@ -0,0 +1,63 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.tensor.functions;
+
+import com.yahoo.tensor.evaluation.Name;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Convenience for cosine similarity between vectors.
+ * cosine_similarity(a, b, mydim) == sum(a*b, mydim) / sqrt(sum(a*a, mydim) * sum(b*b, mydim))
+ * @author arnej
+ */
+public class CosineSimilarity<NAMETYPE extends Name> extends CompositeTensorFunction<NAMETYPE> {
+
+ private final TensorFunction<NAMETYPE> arg1;
+ private final TensorFunction<NAMETYPE> arg2;
+ private final String dimension;
+
+ public CosineSimilarity(TensorFunction<NAMETYPE> argument1,
+ TensorFunction<NAMETYPE> argument2,
+ String dimension)
+ {
+ this.arg1 = argument1;
+ this.arg2 = argument2;
+ this.dimension = dimension;
+ }
+
+ @Override
+ public List<TensorFunction<NAMETYPE>> arguments() { return List.of(arg1, arg2); }
+
+ @Override
+ public TensorFunction<NAMETYPE> withArguments(List<TensorFunction<NAMETYPE>> arguments) {
+ if ( arguments.size() != 2)
+ throw new IllegalArgumentException("CosineSimilarity must have 2 arguments, got " + arguments.size());
+ return new CosineSimilarity<>(arguments.get(0), arguments.get(1), dimension);
+ }
+
+ @Override
+ public PrimitiveTensorFunction<NAMETYPE> toPrimitive() {
+ TensorFunction<NAMETYPE> a = arg1.toPrimitive();
+ TensorFunction<NAMETYPE> b = arg2.toPrimitive();
+ var aa = new Join<>(a, a, ScalarFunctions.multiply());
+ var ab = new Join<>(a, b, ScalarFunctions.multiply());
+ var bb = new Join<>(b, b, ScalarFunctions.multiply());
+ var dot_aa = new Reduce<>(aa, Reduce.Aggregator.sum, dimension);
+ var dot_ab = new Reduce<>(ab, Reduce.Aggregator.sum, dimension);
+ var dot_bb = new Reduce<>(bb, Reduce.Aggregator.sum, dimension);
+ var aabb = new Join<>(dot_aa, dot_bb, ScalarFunctions.multiply());
+ var sqrt_aabb = new Map<>(aabb, ScalarFunctions.sqrt());
+ return new Join<>(dot_ab, sqrt_aabb, ScalarFunctions.divide());
+ }
+
+ @Override
+ public String toString(ToStringContext<NAMETYPE> context) {
+ return "cosine_similarity(" + arg1.toString(context) + ", " + arg2.toString(context) + ", " + dimension + ")";
+ }
+
+ @Override
+ public int hashCode() { return Objects.hash("cosine_similarity", arg1, arg2, dimension); }
+
+}
diff --git a/vespajlib/src/test/java/com/yahoo/tensor/functions/CosineSimilarityTestCase.java b/vespajlib/src/test/java/com/yahoo/tensor/functions/CosineSimilarityTestCase.java
new file mode 100644
index 00000000000..b303e2c1739
--- /dev/null
+++ b/vespajlib/src/test/java/com/yahoo/tensor/functions/CosineSimilarityTestCase.java
@@ -0,0 +1,66 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.tensor.functions;
+
+import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.TensorType;
+import com.yahoo.tensor.evaluation.VariableTensor;
+import org.junit.Test;
+
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author arnej
+ */
+public class CosineSimilarityTestCase {
+
+ @Test
+ public void testVectorSimilarity() {
+ var a = Tensor.from("tensor(x[3]):[ 2.0, 3.0, 6.0]");
+ var b = Tensor.from("tensor(x[3]):[-2.0, 0.0, 0.0]");
+ var c = Tensor.from("tensor(x[3]):[ 0.0, 4.0, 3.0]");
+ var op = new CosineSimilarity<>(new ConstantTensor<>(a), new ConstantTensor<>(b), "x");
+ Tensor result = op.evaluate();
+ assertEquals((-2.0 / 7.0), result.asDouble(), 0.000001);
+ op = new CosineSimilarity<>(new ConstantTensor<>(b), new ConstantTensor<>(a), "x");
+ result = op.evaluate();
+ assertEquals((-2.0 / 7.0), result.asDouble(), 0.000001);
+ op = new CosineSimilarity<>(new ConstantTensor<>(a), new ConstantTensor<>(c), "x");
+ result = op.evaluate();
+ assertEquals((30.0 / 35.0), result.asDouble(), 0.000001);
+ op = new CosineSimilarity<>(new ConstantTensor<>(b), new ConstantTensor<>(c), "x");
+ result = op.evaluate();
+ assertEquals(0.0, result.asDouble(), 0.000001);
+ }
+
+ @Test
+ public void testSimilarityInMixed() {
+ var a = Tensor.from("tensor(c{},yy[3]):{foo:[3.0, 4.0, 0.0],bar:[0.0, -4.0, 3.0]}");
+ var b = Tensor.from("tensor(c{},yy[3]):{foo:[0.0, 4.0, -3.0],bar:[4.0, 0.0, -3.0]}");
+ var op = new CosineSimilarity<>(new ConstantTensor<>(a), new ConstantTensor<>(b), "yy");
+ Tensor result = op.evaluate();
+ var expect = Tensor.from("tensor(c{}):{foo:0.64,bar:-0.36}");
+ assertEquals(expect, result);
+ }
+
+ @Test
+ public void testExpansion() {
+ var tType = TensorType.fromSpec("tensor(vecdim[128])");
+ var a = new VariableTensor<>("left", tType);
+ var b = new VariableTensor<>("right", tType);
+ var op = new CosineSimilarity<>(a, b, "vecdim");
+ assertEquals("join(" +
+ ( "reduce(join(left, right, f(a,b)(a * b)), sum, vecdim), " +
+ "map(" +
+ ( "join(" +
+ ( "reduce(join(left, left, f(a,b)(a * b)), sum, vecdim), " +
+ "reduce(join(right, right, f(a,b)(a * b)), sum, vecdim), " +
+ "f(a,b)(a * b)), " ) +
+ "f(a)(sqrt(a))), " ) +
+ "f(a,b)(a / b)" ) +
+ ")",
+ op.toPrimitive().toString());
+ }
+
+}