Use inference API for embeddings in huggingchat prod (#1037)

* Let the user use HF_TOKEN as auth bearer token in TEI endpoint

* Use inference API for embeddings in huggingchat prod

.env.template

@@ -232,6 +232,18 @@ MODELS=`[
     }
 ]`
 
+TEXT_EMBEDDING_MODELS = `[
+  {
+    "name": "BAAI/bge-small-en-v1.5",
+    "id": "BAAI/bge-small-en-v1.5",
+    "chunkCharLength": 2048,
+    "endpoints": [
+      { "type": "hfapi" }
+    ]
+  }
+]`
+
+
 OLD_MODELS=`[
   {"name":"bigcode/starcoder"},
   {"name":"OpenAssistant/oasst-sft-6-llama-30b-xor"},

src/lib/server/embeddingEndpoints/embeddingEndpoints.ts

@@ -11,6 +11,7 @@ import {
 	embeddingEndpointOpenAI,
 	embeddingEndpointOpenAIParametersSchema,
 } from "./openai/embeddingEndpoints";
+import { embeddingEndpointHfApi, embeddingEndpointHfApiSchema } from "./hfApi/embeddingHfApi";
 
 // parameters passed when generating text
 interface EmbeddingEndpointParameters {
@@ -26,6 +27,7 @@ export const embeddingEndpointSchema = z.discriminatedUnion("type", [
 	embeddingEndpointTeiParametersSchema,
 	embeddingEndpointTransformersJSParametersSchema,
 	embeddingEndpointOpenAIParametersSchema,
+	embeddingEndpointHfApiSchema,
 ]);
 
 type EmbeddingEndpointTypeOptions = z.infer<typeof embeddingEndpointSchema>["type"];
@@ -42,6 +44,7 @@ export const embeddingEndpoints: {
 	tei: embeddingEndpointTei,
 	transformersjs: embeddingEndpointTransformersJS,
 	openai: embeddingEndpointOpenAI,
+	hfapi: embeddingEndpointHfApi,
 };
 
 export default embeddingEndpoints;

src/lib/server/embeddingEndpoints/hfApi/embeddingHfApi.ts

@@ -0,0 +1,51 @@
+import { z } from "zod";
+import type { EmbeddingEndpoint, Embedding } from "../embeddingEndpoints";
+import { chunk } from "$lib/utils/chunk";
+import { HF_TOKEN } from "$env/static/private";
+
+export const embeddingEndpointHfApiSchema = z.object({
+	weight: z.number().int().positive().default(1),
+	model: z.any(),
+	type: z.literal("hfapi"),
+	authorization: z
+		.string()
+		.optional()
+		.transform((v) => (!v && HF_TOKEN ? "Bearer " + HF_TOKEN : v)), // if the header is not set but HF_TOKEN is, use it as the authorization header
+});
+
+export async function embeddingEndpointHfApi(
+	input: z.input<typeof embeddingEndpointHfApiSchema>
+): Promise<EmbeddingEndpoint> {
+	const { model, authorization } = embeddingEndpointHfApiSchema.parse(input);
+	const url = "https://api-inference.huggingface.co/models/" + model.id;
+
+	return async ({ inputs }) => {
+		const batchesInputs = chunk(inputs, 128);
+
+		const batchesResults = await Promise.all(
+			batchesInputs.map(async (batchInputs) => {
+				const response = await fetch(`${url}`, {
+					method: "POST",
+					headers: {
+						Accept: "application/json",
+						"Content-Type": "application/json",
+						...(authorization ? { Authorization: authorization } : {}),
+					},
+					body: JSON.stringify({ inputs: batchInputs }),
+				});
+
+				if (!response.ok) {
+					console.log(await response.text());
+					console.error("Failed to get embeddings from Hugging Face API", response);
+				}
+
+				const embeddings: Embedding[] = await response.json();
+				return embeddings;
+			})
+		);
+
+		const flatAllEmbeddings = batchesResults.flat();
+
+		return flatAllEmbeddings;
+	};
+}

src/lib/server/embeddingEndpoints/tei/embeddingEndpoints.ts

@@ -1,13 +1,17 @@
 import { z } from "zod";
 import type { EmbeddingEndpoint, Embedding } from "../embeddingEndpoints";
 import { chunk } from "$lib/utils/chunk";
+import { HF_TOKEN } from "$env/static/private";
 
 export const embeddingEndpointTeiParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
 	model: z.any(),
 	type: z.literal("tei"),
 	url: z.string().url(),
-	authorization: z.string().optional(),
+	authorization: z
+		.string()
+		.optional()
+		.transform((v) => (!v && HF_TOKEN ? "Bearer " + HF_TOKEN : v)), // if the header is not set but HF_TOKEN is, use it as the authorization header
 });
 
 const getModelInfoByUrl = async (url: string, authorization?: string) => {

src/lib/server/embeddingModels.ts

@@ -73,6 +73,10 @@ const addEndpoint = (m: Awaited<ReturnType<typeof processEmbeddingModel>>) => ({
 						return embeddingEndpoints.transformersjs(args);
 					case "openai":
 						return embeddingEndpoints.openai(args);
+					case "hfapi":
+						return embeddingEndpoints.hfapi(args);
+					default:
+						throw new Error(`Unknown endpoint type: ${args}`);
 				}
 			}