diff options
author | Jon Bratseth <bratseth@gmail.com> | 2023-06-26 17:23:35 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-06-26 17:23:35 +0200 |
commit | 55fb18fa6800ddfc8abc33d6bad602710a5c91ae (patch) | |
tree | 4f5ba72143b874153aa65560c2c172f6fe9d7ff3 | |
parent | 0361837883a3b5c02db3c215880ce90ba3ad8c65 (diff) | |
parent | 89150530a47690fa0df603069789002f79ae7123 (diff) |
Merge pull request #27555 from vespa-engine/arnej/add-cosine-similarity
Arnej/add cosine similarity
6 files changed, 228 insertions, 2 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..76d007dd633 100644 --- a/vespajlib/abi-spec.json +++ b/vespajlib/abi-spec.json @@ -1705,6 +1705,24 @@ ], "fields" : [ ] }, + "com.yahoo.tensor.functions.CosineSimilarity" : { + "superClass" : "com.yahoo.tensor.functions.TensorFunction", + "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.TensorType type(com.yahoo.tensor.evaluation.TypeContext)", + "public com.yahoo.tensor.Tensor evaluate(com.yahoo.tensor.evaluation.EvaluationContext)", + "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" : [ ], @@ -1741,7 +1759,7 @@ "fields" : [ ] }, "com.yahoo.tensor.functions.EuclideanDistance" : { - "superClass" : "com.yahoo.tensor.functions.CompositeTensorFunction", + "superClass" : "com.yahoo.tensor.functions.TensorFunction", "interfaces" : [ ], "attributes" : [ "public" @@ -1750,6 +1768,8 @@ "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.TensorType type(com.yahoo.tensor.evaluation.TypeContext)", + "public com.yahoo.tensor.Tensor evaluate(com.yahoo.tensor.evaluation.EvaluationContext)", "public com.yahoo.tensor.functions.PrimitiveTensorFunction toPrimitive()", "public java.lang.String toString(com.yahoo.tensor.functions.ToStringContext)", "public int hashCode()" 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..ebb8a11fd8a --- /dev/null +++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/CosineSimilarity.java @@ -0,0 +1,93 @@ +// 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.EvaluationContext; +import com.yahoo.tensor.evaluation.Name; +import com.yahoo.tensor.evaluation.TypeContext; +import com.yahoo.tensor.Tensor; +import com.yahoo.tensor.TensorType; +import com.yahoo.tensor.TensorType.Dimension; + +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 TensorFunction<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 TensorType type(TypeContext<NAMETYPE> context) { + TensorType t1 = arg1.toPrimitive().type(context); + TensorType t2 = arg2.toPrimitive().type(context); + var d1 = t1.dimension(dimension); + var d2 = t2.dimension(dimension); + if (d1.isEmpty() || d2.isEmpty() + || d1.get().type() != Dimension.Type.indexedBound + || d2.get().type() != Dimension.Type.indexedBound + || d1.get().size().get() != d2.get().size().get()) + { + throw new IllegalArgumentException("cosine_similarity expects both arguments to have the '" + + dimension + "' dimension with same size, but input types were " + + t1 + " and " + t2); + } + // Finds the type this produces by first converting it to a primitive function + return toPrimitive().type(context); + } + + /** Evaluates this by first converting it to a primitive function */ + @Override + public Tensor evaluate(EvaluationContext<NAMETYPE> context) { + return toPrimitive().evaluate(context); + } + + @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/main/java/com/yahoo/tensor/functions/EuclideanDistance.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/EuclideanDistance.java index 25399416c29..f9fc8e195d3 100644 --- a/vespajlib/src/main/java/com/yahoo/tensor/functions/EuclideanDistance.java +++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/EuclideanDistance.java @@ -1,7 +1,12 @@ // 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.EvaluationContext; import com.yahoo.tensor.evaluation.Name; +import com.yahoo.tensor.evaluation.TypeContext; +import com.yahoo.tensor.Tensor; +import com.yahoo.tensor.TensorType; +import com.yahoo.tensor.TensorType.Dimension; import java.util.Collections; import java.util.List; @@ -12,7 +17,7 @@ import java.util.Objects; * euclidean_distance(a, b, mydim) == sqrt(sum(pow(a-b, 2), mydim)) * @author arnej */ -public class EuclideanDistance<NAMETYPE extends Name> extends CompositeTensorFunction<NAMETYPE> { +public class EuclideanDistance<NAMETYPE extends Name> extends TensorFunction<NAMETYPE> { private final TensorFunction<NAMETYPE> arg1; private final TensorFunction<NAMETYPE> arg2; @@ -38,6 +43,31 @@ public class EuclideanDistance<NAMETYPE extends Name> extends CompositeTensorFun } @Override + public TensorType type(TypeContext<NAMETYPE> context) { + TensorType t1 = arg1.toPrimitive().type(context); + TensorType t2 = arg2.toPrimitive().type(context); + var d1 = t1.dimension(dimension); + var d2 = t2.dimension(dimension); + if (d1.isEmpty() || d2.isEmpty() + || d1.get().type() != Dimension.Type.indexedBound + || d2.get().type() != Dimension.Type.indexedBound + || d1.get().size().get() != d2.get().size().get()) + { + throw new IllegalArgumentException("euclidean_distance expects both arguments to have the '" + + dimension + "' dimension with same size, but input types were " + + t1 + " and " + t2); + } + // Finds the type this produces by first converting it to a primitive function + return toPrimitive().type(context); + } + + /** Evaluates this by first converting it to a primitive function */ + @Override + public Tensor evaluate(EvaluationContext<NAMETYPE> context) { + return toPrimitive().evaluate(context); + } + + @Override public PrimitiveTensorFunction<NAMETYPE> toPrimitive() { TensorFunction<NAMETYPE> primitive1 = arg1.toPrimitive(); TensorFunction<NAMETYPE> primitive2 = arg2.toPrimitive(); 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()); + } + +} |