Allow users to specify similarity

jonathan-buttner · Mar 19, 2024 · e7cb97d · e7cb97d
1 parent edc45f5
commit e7cb97d
Show file tree

Hide file tree

Showing 17 changed files with 403 additions and 43 deletions.
diff --git a/...java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsAction.java b/...java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsAction.java
@@ -31,7 +31,7 @@ public CohereEmbeddingsAction(Sender sender, CohereEmbeddingsModel model) {
         Objects.requireNonNull(model);
         this.sender = Objects.requireNonNull(sender);
         this.failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(
-            model.getServiceSettings().getCommonSettings().getUri(),
+            model.getServiceSettings().getCommonSettings().uri(),
             "Cohere embeddings"
         );
         requestCreator = new CohereEmbeddingsExecutableRequestCreator(model);

diff --git a/...search/xpack/inference/external/http/sender/CohereEmbeddingsExecutableRequestCreator.java b/...search/xpack/inference/external/http/sender/CohereEmbeddingsExecutableRequestCreator.java
@@ -37,7 +37,7 @@ private static ResponseHandler createEmbeddingsHandler() {
 
     public CohereEmbeddingsExecutableRequestCreator(CohereEmbeddingsModel model) {
         this.model = Objects.requireNonNull(model);
-        account = new CohereAccount(this.model.getServiceSettings().getCommonSettings().getUri(), this.model.getSecretSettings().apiKey());
+        account = new CohereAccount(this.model.getServiceSettings().getCommonSettings().uri(), this.model.getSecretSettings().apiKey());
     }
 
     @Override

diff --git a/...va/org/elasticsearch/xpack/inference/external/request/cohere/CohereEmbeddingsRequest.java b/...va/org/elasticsearch/xpack/inference/external/request/cohere/CohereEmbeddingsRequest.java
@@ -46,7 +46,7 @@ public CohereEmbeddingsRequest(CohereAccount account, List<String> input, Cohere
         this.input = Objects.requireNonNull(input);
         uri = buildUri(this.account.url(), "Cohere", CohereEmbeddingsRequest::buildDefaultUri);
         taskSettings = embeddingsModel.getTaskSettings();
-        model = embeddingsModel.getServiceSettings().getCommonSettings().getModelId();
+        model = embeddingsModel.getServiceSettings().getCommonSettings().modelId();
         embeddingType = embeddingsModel.getServiceSettings().getEmbeddingType();
         inferenceEntityId = embeddingsModel.getInferenceEntityId();
     }

diff --git a/...erence/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java b/...erence/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java
@@ -216,13 +216,16 @@ public void checkModelConfig(Model model, ActionListener<Model> listener) {
     }
 
     private CohereEmbeddingsModel updateModelWithEmbeddingDetails(CohereEmbeddingsModel model, int embeddingSize) {
+        var similarityFromModel = model.getServiceSettings().similarity();
+        var similarityToUse = similarityFromModel == null ? SimilarityMeasure.DOT_PRODUCT : similarityFromModel;
+
         CohereEmbeddingsServiceSettings serviceSettings = new CohereEmbeddingsServiceSettings(
             new CohereServiceSettings(
-                model.getServiceSettings().getCommonSettings().getUri(),
-                SimilarityMeasure.DOT_PRODUCT,
+                model.getServiceSettings().getCommonSettings().uri(),
+                similarityToUse,
                 embeddingSize,
-                model.getServiceSettings().getCommonSettings().getMaxInputTokens(),
-                model.getServiceSettings().getCommonSettings().getModelId()
+                model.getServiceSettings().getCommonSettings().maxInputTokens(),
+                model.getServiceSettings().getCommonSettings().modelId()
             ),
             model.getServiceSettings().getEmbeddingType()
         );

diff --git a/...rc/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java b/...rc/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java
@@ -65,10 +65,10 @@ public static CohereServiceSettings fromMap(Map<String, Object> map, Configurati
             throw validationException;
         }
 
-        return new CohereServiceSettings(uri, similarity, dims, maxInputTokens, getModelId(oldModelId, modelId));
+        return new CohereServiceSettings(uri, similarity, dims, maxInputTokens, modelId(oldModelId, modelId));
     }
 
-    private static String getModelId(@Nullable String model, @Nullable String modelId) {
+    private static String modelId(@Nullable String model, @Nullable String modelId) {
         return modelId != null ? modelId : model;
     }
 
@@ -110,23 +110,25 @@ public CohereServiceSettings(StreamInput in) throws IOException {
         modelId = in.readOptionalString();
     }
 
-    public URI getUri() {
+    public URI uri() {
         return uri;
     }
 
-    public SimilarityMeasure getSimilarity() {
+    @Override
+    public SimilarityMeasure similarity() {
         return similarity;
     }
 
-    public Integer getDimensions() {
+    @Override
+    public Integer dimensions() {
         return dimensions;
     }
 
-    public Integer getMaxInputTokens() {
+    public Integer maxInputTokens() {
         return maxInputTokens;
     }
 
-    public String getModelId() {
+    public String modelId() {
         return modelId;
     }
 

diff --git a/...ticsearch/xpack/inference/services/cohere/embeddings/CohereEmbeddingsServiceSettings.java b/...ticsearch/xpack/inference/services/cohere/embeddings/CohereEmbeddingsServiceSettings.java
@@ -15,6 +15,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ServiceSettings;
+import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
@@ -93,6 +94,16 @@ public CohereServiceSettings getCommonSettings() {
         return commonSettings;
     }
 
+    @Override
+    public SimilarityMeasure similarity() {
+        return commonSettings.similarity();
+    }
+
+    @Override
+    public Integer dimensions() {
+        return commonSettings.dimensions();
+    }
+
     public CohereEmbeddingType getEmbeddingType() {
         return embeddingType;
     }

diff --git a/.../main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceService.java b/.../main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceService.java
@@ -61,7 +61,7 @@ public void checkModelConfig(Model model, ActionListener<Model> listener) {
     private static HuggingFaceEmbeddingsModel updateModelWithEmbeddingDetails(HuggingFaceEmbeddingsModel model, int embeddingSize) {
         var serviceSettings = new HuggingFaceServiceSettings(
             model.getServiceSettings().uri(),
-            null, // Similarity measure is unknown
+            model.getServiceSettings().similarity(), // we don't know the similarity but use whatever the user specified
             embeddingSize,
             model.getTokenLimit()
         );

diff --git a/...va/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceSettings.java b/...va/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceSettings.java
@@ -144,10 +144,12 @@ public URI uri() {
         return uri;
     }
 
+    @Override
     public SimilarityMeasure similarity() {
         return similarity;
     }
 
+    @Override
     public Integer dimensions() {
         return dimensions;
     }

diff --git a/...erence/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java b/...erence/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java
@@ -239,11 +239,14 @@ private OpenAiEmbeddingsModel updateModelWithEmbeddingDetails(OpenAiEmbeddingsMo
             );
         }
 
+        var similarityFromModel = model.getServiceSettings().similarity();
+        var similarityToUse = similarityFromModel == null ? SimilarityMeasure.DOT_PRODUCT : similarityFromModel;
+
         OpenAiEmbeddingsServiceSettings serviceSettings = new OpenAiEmbeddingsServiceSettings(
             model.getServiceSettings().modelId(),
             model.getServiceSettings().uri(),
             model.getServiceSettings().organizationId(),
-            SimilarityMeasure.DOT_PRODUCT,
+            similarityToUse,
             embeddingSize,
             model.getServiceSettings().maxInputTokens(),
             model.getServiceSettings().dimensionsSetByUser()

diff --git a/...ticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsServiceSettings.java b/...ticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsServiceSettings.java
@@ -191,10 +191,12 @@ public String organizationId() {
         return organizationId;
     }
 
+    @Override
     public SimilarityMeasure similarity() {
         return similarity;
     }
 
+    @Override
     public Integer dimensions() {
         return dimensions;
     }

diff --git a/...g/elasticsearch/xpack/inference/external/request/cohere/CohereEmbeddingsRequestTests.java b/...g/elasticsearch/xpack/inference/external/request/cohere/CohereEmbeddingsRequestTests.java
@@ -179,7 +179,7 @@ public void testCreateRequest_TruncateNone() throws URISyntaxException, IOExcept
     }
 
     public static CohereEmbeddingsRequest createRequest(List<String> input, CohereEmbeddingsModel model) throws URISyntaxException {
-        var account = new CohereAccount(model.getServiceSettings().getCommonSettings().getUri(), model.getSecretSettings().apiKey());
+        var account = new CohereAccount(model.getServiceSettings().getCommonSettings().uri(), model.getSecretSettings().apiKey());
         return new CohereEmbeddingsRequest(account, input, model);
     }
 }
diff --git a/...st/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettingsTests.java b/...st/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettingsTests.java
@@ -151,7 +151,7 @@ public void testFromMap_PrefersModelId_OverModel() {
 
     public void testFromMap_MissingUrl_DoesNotThrowException() {
         var serviceSettings = CohereServiceSettings.fromMap(new HashMap<>(Map.of()), ConfigurationParseContext.PERSISTENT);
-        assertNull(serviceSettings.getUri());
+        assertNull(serviceSettings.uri());
     }
 
     public void testFromMap_EmptyUrl_ThrowsError() {