From 2ef62a93bb8316b837d744a7895d3bb5df5171ef Mon Sep 17 00:00:00 2001
From: Tasmiyah Iqbal <tasmay@users.noreply.github.com>
Date: Tue, 25 Jun 2024 14:16:41 -0400
Subject: [PATCH] feat: added support for embeddings via HuggingFace Inference
 API (#929)

---
 .changeset/funny-pugs-hide.md                 |  6 +++
 examples/huggingface/embeddingApi.ts          | 51 ++++++++++++++++++
 .../src/embeddings/HuggingFaceEmbedding.ts    | 53 +++++++++++++++++++
 packages/llamaindex/src/embeddings/index.ts   |  1 +
 4 files changed, 111 insertions(+)
 create mode 100644 .changeset/funny-pugs-hide.md
 create mode 100644 examples/huggingface/embeddingApi.ts

diff --git a/.changeset/funny-pugs-hide.md b/.changeset/funny-pugs-hide.md
new file mode 100644
index 000000000..992e553a3
--- /dev/null
+++ b/.changeset/funny-pugs-hide.md
@@ -0,0 +1,6 @@
+---
+"llamaindex": patch
+"@llamaindex/examples": patch
+---
+
+feat: added support for embeddings via HuggingFace Inference API
diff --git a/examples/huggingface/embeddingApi.ts b/examples/huggingface/embeddingApi.ts
new file mode 100644
index 000000000..a89df2703
--- /dev/null
+++ b/examples/huggingface/embeddingApi.ts
@@ -0,0 +1,51 @@
+import fs from "node:fs/promises";
+
+import {
+  Document,
+  HuggingFaceInferenceAPI,
+  HuggingFaceInferenceAPIEmbedding,
+  Settings,
+  VectorStoreIndex,
+} from "llamaindex";
+
+if (!process.env.HUGGING_FACE_TOKEN) {
+  throw new Error("Please set the HUGGING_FACE_TOKEN environment variable.");
+}
+
+// Update embed model with HuggingFaceAPIEmbedding
+Settings.embedModel = new HuggingFaceInferenceAPIEmbedding({
+  model: "mixedbread-ai/mxbai-embed-large-v1",
+  accessToken: process.env.HUGGING_FACE_TOKEN,
+});
+// Omit this if you want to use OpenAI LLM as default otherwise set to your preferred LLM
+Settings.llm = new HuggingFaceInferenceAPI({
+  model: "meta-llama/Meta-Llama-3-8B-Instruct",
+  accessToken: process.env.HUGGING_FACE_TOKEN,
+});
+
+async function main() {
+  // Load essay from abramov.txt in Node
+  const path = "node_modules/llamaindex/examples/abramov.txt";
+
+  const essay = await fs.readFile(path, "utf-8");
+
+  // Create Document object with essay
+  const document = new Document({ text: essay, id_: path });
+
+  // Split text and create embeddings. Store them in a VectorStoreIndex
+  const index = await VectorStoreIndex.fromDocuments([document]);
+
+  // Query the index
+  const queryEngine = index.asQueryEngine();
+  const stream = await queryEngine.query({
+    query: "What did the author do in college?",
+    stream: true,
+  });
+
+  // Output response
+  for await (const chunk of stream) {
+    process.stdout.write(chunk.response);
+  }
+}
+
+main().catch(console.error);
diff --git a/packages/llamaindex/src/embeddings/HuggingFaceEmbedding.ts b/packages/llamaindex/src/embeddings/HuggingFaceEmbedding.ts
index dce8b0208..59771d3d3 100644
--- a/packages/llamaindex/src/embeddings/HuggingFaceEmbedding.ts
+++ b/packages/llamaindex/src/embeddings/HuggingFaceEmbedding.ts
@@ -1,3 +1,4 @@
+import { HfInference } from "@huggingface/inference";
 import { lazyLoadTransformers } from "../internal/deps/transformers.js";
 import { BaseEmbedding } from "./types.js";
 
@@ -46,3 +47,55 @@ export class HuggingFaceEmbedding extends BaseEmbedding {
     return Array.from(output.data);
   }
 }
+
+// Workaround to get the Options type from @huggingface/inference@2.7.0
+type HfInferenceOptions = ConstructorParameters<typeof HfInference>[1];
+
+export type HFConfig = HfInferenceOptions & {
+  model: string;
+  accessToken: string;
+  endpoint?: string;
+};
+
+/**
+ * Uses feature extraction from Hugging Face's Inference API to generate embeddings.
+ *
+ * Set the `model` and `accessToken` parameter in the constructor, e.g.:
+ * ```
+ * new HuggingFaceInferenceAPIEmbedding({
+ *     model: HuggingFaceEmbeddingModelType.XENOVA_ALL_MPNET_BASE_V2,
+ *     accessToken: "<your-access-token>"
+ * });
+ * ```
+ *
+ * @extends BaseEmbedding
+ */
+export class HuggingFaceInferenceAPIEmbedding extends BaseEmbedding {
+  model: string;
+  hf: HfInference;
+
+  constructor(init: HFConfig) {
+    super();
+    const { model, accessToken, endpoint, ...hfInferenceOpts } = init;
+
+    this.hf = new HfInference(accessToken, hfInferenceOpts);
+    this.model = model;
+    if (endpoint) this.hf.endpoint(endpoint);
+  }
+
+  async getTextEmbedding(text: string): Promise<number[]> {
+    const res = await this.hf.featureExtraction({
+      model: this.model,
+      inputs: text,
+    });
+    return res as number[];
+  }
+
+  async getTextEmbeddings(texts: string[]): Promise<Array<number[]>> {
+    const res = await this.hf.featureExtraction({
+      model: this.model,
+      inputs: texts,
+    });
+    return res as number[][];
+  }
+}
diff --git a/packages/llamaindex/src/embeddings/index.ts b/packages/llamaindex/src/embeddings/index.ts
index 012ce5307..59ffb661d 100644
--- a/packages/llamaindex/src/embeddings/index.ts
+++ b/packages/llamaindex/src/embeddings/index.ts
@@ -1,6 +1,7 @@
 export { DeepInfraEmbedding } from "./DeepInfraEmbedding.js";
 export { FireworksEmbedding } from "./fireworks.js";
 export * from "./GeminiEmbedding.js";
+export { HuggingFaceInferenceAPIEmbedding } from "./HuggingFaceEmbedding.js";
 export * from "./JinaAIEmbedding.js";
 export * from "./MistralAIEmbedding.js";
 export * from "./MultiModalEmbedding.js";
-- 
GitLab