feat: added support for embeddings via HuggingFace Inference API (#929)

.changeset/funny-pugs-hide.md

@@ -0,0 +1,6 @@
+---
+"llamaindex": patch
+"@llamaindex/examples": patch
+---
+
+feat: added support for embeddings via HuggingFace Inference API

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);

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[][];
+  }
+}

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";