feat: HFEI support

Author: tazarov

src/main/java/tech/amikos/chromadb/embeddings/hf/HuggingFaceEmbeddingFunction.java

@@ -17,12 +17,15 @@
 public class HuggingFaceEmbeddingFunction implements EmbeddingFunction {
     public static final String DEFAULT_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2";
     public static final String DEFAULT_BASE_API = "https://api-inference.huggingface.co/pipeline/feature-extraction/";
+    public static final String HFEI_API_PATH = "/embed";
     public static final String HF_API_KEY_ENV = "HF_API_KEY";
+    public static final String API_TYPE_CONFIG_KEY = "apiType";
     private final OkHttpClient client = new OkHttpClient();
     private final Map<String, Object> configParams = new HashMap<>();
     private static final Gson gson = new Gson();
 
     private static final List<WithParam> defaults = Arrays.asList(
+            new WithAPIType(APIType.HF_API),
             WithParam.baseAPI(DEFAULT_BASE_API),
             WithParam.defaultModel(DEFAULT_MODEL_NAME)
     );
@@ -46,14 +49,21 @@ public HuggingFaceEmbeddingFunction(WithParam... params) throws EFException {
     }
 
     public CreateEmbeddingResponse createEmbedding(CreateEmbeddingRequest req) throws EFException {
-        Request request = new Request.Builder()
-                .url(this.configParams.get(Constants.EF_PARAMS_BASE_API).toString() + this.configParams.get(Constants.EF_PARAMS_MODEL).toString())
+        Request.Builder rb = new Request.Builder()
+
                 .post(RequestBody.create(req.json(), JSON))
                 .addHeader("Accept", "application/json")
                 .addHeader("Content-Type", "application/json")
-                .addHeader("User-Agent", Constants.HTTP_AGENT)
-                .addHeader("Authorization", "Bearer " + configParams.get(Constants.EF_PARAMS_API_KEY).toString())
-                .build();
+                .addHeader("User-Agent", Constants.HTTP_AGENT);
+        if (configParams.containsKey(API_TYPE_CONFIG_KEY) && configParams.get(API_TYPE_CONFIG_KEY).equals(APIType.HFEI_API)) {
+            rb.url(this.configParams.get(Constants.EF_PARAMS_BASE_API).toString() + HFEI_API_PATH);
+        } else {
+            rb.url(this.configParams.get(Constants.EF_PARAMS_BASE_API).toString() + this.configParams.get(Constants.EF_PARAMS_MODEL).toString());
+        }
+        if (configParams.containsKey(Constants.EF_PARAMS_API_KEY)) {
+            rb.addHeader("Authorization", "Bearer " + configParams.get(Constants.EF_PARAMS_API_KEY).toString());
+        }
+        Request request = rb.build();
         try (Response response = client.newCall(request).execute()) {
             if (!response.isSuccessful()) {
                 throw new IOException("Unexpected code " + response);
@@ -86,4 +96,22 @@ public List<Embedding> embedDocuments(String[] documents) throws EFException {
         CreateEmbeddingResponse response = this.createEmbedding(new CreateEmbeddingRequest().inputs(documents));
         return response.getEmbeddings().stream().map(Embedding::fromList).collect(Collectors.toList());
     }
+
+    public static class WithAPIType extends WithParam {
+        private final APIType apiType;
+
+        public WithAPIType(APIType apitype) {
+            this.apiType = apitype;
+        }
+
+        @Override
+        public void apply(Map<String, Object> params) {
+            params.put(API_TYPE_CONFIG_KEY, apiType);
+        }
+    }
+
+    public static enum APIType{
+        HF_API,
+        HFEI_API
+    }
 }

src/test/java/tech/amikos/chromadb/embeddings/hf/TestHuggingFaceEmbeddings.java

@@ -2,6 +2,8 @@
 
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.containers.wait.strategy.Wait;
 import tech.amikos.chromadb.*;
 import tech.amikos.chromadb.embeddings.EmbeddingFunction;
 import tech.amikos.chromadb.embeddings.WithParam;
@@ -13,10 +15,23 @@
 import static org.junit.Assert.assertNotNull;
 
 public class TestHuggingFaceEmbeddings {
+    static GenericContainer hfeiContainer;
 
     @BeforeClass
-    public static void setup() {
+    public static void setup() throws Exception {
         Utils.loadEnvFile(".env");
+
+        try {
+            hfeiContainer = new GenericContainer("ghcr.io/huggingface/text-embeddings-inference:cpu-1.5.0")
+                    .withCommand("--model-id Snowflake/snowflake-arctic-embed-s --revision main")
+                    .withExposedPorts(80)
+                    .waitingFor(Wait.forHttp("/").forStatusCode(200));
+            hfeiContainer.start();
+            System.setProperty("HFEI_URL", "http://" + hfeiContainer.getHost() + ":" + hfeiContainer.getMappedPort(80));
+        } catch (Exception e) {
+            System.err.println("HFEI container failed to start");
+            throw e;
+        }
     }
 
     @Test
@@ -45,5 +60,15 @@ public void testWithModel() throws ApiException, EFException {
         assertNotNull(results);
         assertEquals(768, results.getDimensions());
     }
+
+    @Test
+    public void testWithURL() throws EFException {
+        EmbeddingFunction ef = new HuggingFaceEmbeddingFunction(
+                WithParam.baseAPI(System.getProperty("HFEI_URL")),
+                new HuggingFaceEmbeddingFunction.WithAPIType(HuggingFaceEmbeddingFunction.APIType.HFEI_API));
+        Embedding results = ef.embedQuery("How are you?");
+        assertNotNull(results);
+        assertEquals(384, results.getDimensions());
+    }
 }