Add embedding models configurable, from both transformers.js and TEI

.env

@@ -42,6 +42,19 @@ CA_PATH=#
 CLIENT_KEY_PASSWORD=#
 REJECT_UNAUTHORIZED=true
 
+TEXT_EMBEDDING_MODELS = `[
+  {
+    "name": "Xenova/gte-small",
+    "displayName": "Xenova/gte-small",
+    "description": "Local embedding model running on the server.",
+    "maxSequenceLength": 512,
+    "endpoints": [
+      { "type": "xenova" }
+    ]
+  }
+]`
+
+
 # 'name', 'userMessageToken', 'assistantMessageToken' are required
 MODELS=`[
     {

.env.template

@@ -196,6 +196,7 @@ MODELS=`[
       "max_new_tokens" : 8192,
       "stop" : ["</s>"]
     },
+    "embeddingModelName": "thenlper/gte-base",
     "promptExamples" : [
       {
         "title": "Write an email from bullet list",
@@ -215,6 +216,27 @@ OLD_MODELS=`[{"name":"bigcode/starcoder"}, {"name":"OpenAssistant/oasst-sft-6-ll
 
 TASK_MODEL='mistralai/Mistral-7B-Instruct-v0.2'
 
+# Default to using the first text embedding model when not specifying 'embeddingModelName' in the model itself.
+TEXT_EMBEDDING_MODELS = `[
+  {
+    "name": "Xenova/gte-small",
+    "displayName": "Xenova/gte-small",
+    "description": "Local embedding model running on the server.",
+    "maxSequenceLength": 512,
+    "endpoints": [
+      { "type": "xenova" }
+    ]
+  },
+  {
+    "name": "thenlper/gte-base",
+    "displayName": "thenlper/gte-base",
+    "description": "Hosted embedding model running on the cloud somewhere.",
+    "maxSequenceLength": 512,
+    "endpoints": [
+      { "type": "tei", "http://localhost:8080/" } # Make sure the MAX_CLIENT_BATCH_SIZE env in TEI is large enough to accommodate web search chunks.  
+    ]
+  }
+]`
 
 APP_BASE="/chat"
 PUBLIC_ORIGIN=https://huggingface.co

src/lib/server/embeddingEndpoints/embeddingEndpoints.ts

@@ -0,0 +1,33 @@
+import { embeddingEndpointTei, embeddingEndpointTeiParametersSchema } from "./tei/embeddingEndpoints";
+import { z } from "zod";
+import embeddingEndpointXenova, { embeddingEndpointXenovaParametersSchema } from "./xenova/embeddingEndpoints";
+
+// parameters passed when generating text
+interface EmbeddingEndpointParameters {
+	inputs: string[]
+}
+
+interface CommonEmbeddingEndpoint {
+	weight: number;
+}
+
+// type signature for the endpoint
+export type EmbeddingEndpoint = (
+	params: EmbeddingEndpointParameters
+) => Promise<any>; // TODO: type
+
+// generator function that takes in parameters for defining the endpoint and return the endpoint
+export type EmbeddingEndpointGenerator<T extends CommonEmbeddingEndpoint> = (parameters: T) => EmbeddingEndpoint;
+
+// list of all endpoint generators
+export const embeddingEndpoints = {
+	tei: embeddingEndpointTei,
+	xenova: embeddingEndpointXenova
+};
+
+export const embeddingEndpointSchema = z.discriminatedUnion("type", [
+	embeddingEndpointTeiParametersSchema,
+	embeddingEndpointXenovaParametersSchema
+]);
+
+export default embeddingEndpoints;

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

@@ -0,0 +1,31 @@
+import { HF_ACCESS_TOKEN, HF_TOKEN } from "$env/static/private";
+import { featureExtraction } from "@huggingface/inference";
+import { z } from "zod";
+import type { EmbeddingEndpoint } from "../embeddingEndpoints";
+
+export const embeddingEndpointTeiParametersSchema = z.object({
+	weight: z.number().int().positive().default(1),
+	model: z.any(),
+	type: z.literal("tei"),
+	url: z.string().url(),
+});
+
+export function embeddingEndpointTei(input: z.input<typeof embeddingEndpointTeiParametersSchema>): EmbeddingEndpoint {
+	const { url } = embeddingEndpointTeiParametersSchema.parse(input);
+	return async ({ inputs }) => {
+		const { origin } = new URL(url)
+
+		const response = await fetch(`${origin}/embed`, {
+			method: 'POST',
+			headers: {
+				'Accept': 'application/json',
+				'Content-Type': 'application/json'
+			},
+			body: JSON.stringify({ inputs, normalize: true, truncate: true })
+		});
+
+		return response.json();
+	};
+}
+
+export default embeddingEndpointTei;

src/lib/server/embeddingEndpoints/xenova/embeddingEndpoints.ts

@@ -0,0 +1,51 @@
+import { z } from "zod";
+import type { EmbeddingEndpoint } from "../embeddingEndpoints";
+import type { Tensor, Pipeline } from "@xenova/transformers";
+import { pipeline } from "@xenova/transformers";
+
+export const embeddingEndpointXenovaParametersSchema = z.object({
+	weight: z.number().int().positive().default(1),
+	model: z.any(),
+	type: z.literal("xenova"),
+});
+
+
+// Use the Singleton pattern to enable lazy construction of the pipeline.
+class XenovaModelsSingleton {
+	static instances: Array<[string, Promise<Pipeline>]> = [];
+
+	static async getInstance(modelName: string): Promise<Pipeline> {
+		const modelPipeline = this.instances.find(([name]) => name === modelName)
+
+		if (modelPipeline) {
+			return modelPipeline[1];
+		}
+
+		const newModelPipeline = pipeline("feature-extraction", modelName)
+		this.instances.push([modelName, newModelPipeline])
+
+		return newModelPipeline;
+	}
+}
+
+
+export async function calculateEmbedding(
+	modelName: string,
+	inputs: string[]
+) {
+	const extractor = await XenovaModelsSingleton.getInstance(modelName);
+	const output: Tensor = await extractor(inputs, { pooling: "mean", normalize: true });
+
+	return output.tolist();
+}
+
+
+export function embeddingEndpointXenova(input: z.input<typeof embeddingEndpointXenovaParametersSchema>): EmbeddingEndpoint {
+	const { model } = embeddingEndpointXenovaParametersSchema.parse(input);
+
+	return async ({ inputs }) => {
+		return calculateEmbedding(model.name, inputs);
+	};
+}
+
+export default embeddingEndpointXenova;

src/lib/server/embeddingModels.ts

@@ -0,0 +1,74 @@
+import {
+	TEXT_EMBEDDING_MODELS,
+} from "$env/static/private";
+
+import { z } from "zod";
+import { sum } from "$lib/utils/sum";
+import embeddingEndpoints, { embeddingEndpointSchema, type EmbeddingEndpoint } from "./embeddingEndpoints/embeddingEndpoints";
+import embeddingEndpointXenova from "./embeddingEndpoints/xenova/embeddingEndpoints";
+
+const modelConfig = z.object({
+	/** Used as an identifier in DB */
+	id: z.string().optional(),
+	/** Used to link to the model page, and for inference */
+	name: z.string().min(1),
+	displayName: z.string().min(1).optional(),
+	description: z.string().min(1).optional(),
+	websiteUrl: z.string().url().optional(),
+	modelUrl: z.string().url().optional(),
+	endpoints: z.array(embeddingEndpointSchema).optional(),
+	maxSequenceLength: z.number().positive(),
+});
+
+const embeddingModelsRaw = z.array(modelConfig).parse(JSON.parse(TEXT_EMBEDDING_MODELS));
+
+const processEmbeddingModel = async (m: z.infer<typeof modelConfig>) => ({
+	...m,
+	id: m.id || m.name,
+});
+
+const addEndpoint = (m: Awaited<ReturnType<typeof processEmbeddingModel>>) => ({
+	...m,
+	getEndpoint: async (): Promise<EmbeddingEndpoint> => {
+		if (!m.endpoints) {
+			return embeddingEndpointXenova({
+				type: "xenova",
+				weight: 1,
+				model: m,
+			});
+		}
+
+		const totalWeight = sum(m.endpoints.map((e) => e.weight));
+
+		let random = Math.random() * totalWeight;
+
+		for (const endpoint of m.endpoints) {
+			if (random < endpoint.weight) {
+				const args = { ...endpoint, model: m };
+
+				switch (args.type) {
+					case "tei":
+						return embeddingEndpoints.tei(args);
+					case "xenova":
+						return embeddingEndpoints.xenova(args);
+				}
+			}
+
+			random -= endpoint.weight;
+		}
+
+		throw new Error(`Failed to select endpoint`);
+	},
+});
+
+export const embeddingModels = await Promise.all(embeddingModelsRaw.map((e) => processEmbeddingModel(e).then(addEndpoint)));
+
+export const defaultEmbeddingModel = embeddingModels[0];
+
+export const validateEmbeddingModel = (_models: EmbeddingBackendModel[]) => {
+	// Zod enum function requires 2 parameters
+	return z.enum([_models[0].id, ..._models.slice(1).map((m) => m.id)]);
+};
+
+
+export type EmbeddingBackendModel = typeof defaultEmbeddingModel;

src/lib/server/models.ts

@@ -66,6 +66,7 @@ const modelConfig = z.object({
 		.optional(),
 	multimodal: z.boolean().default(false),
 	unlisted: z.boolean().default(false),
+	embeddingModelName: z.string().optional()
 });
 
 const modelsRaw = z.array(modelConfig).parse(JSON.parse(MODELS));

src/lib/server/sentenceSimilarity.ts

@@ -0,0 +1,38 @@
+import { dot } from "@xenova/transformers";
+import type { EmbeddingBackendModel } from "./embeddingModels";
+
+// see here: https://github.com/nmslib/hnswlib/blob/359b2ba87358224963986f709e593d799064ace6/README.md?plain=1#L34
+function innerProduct(embeddingA: number[], embeddingB: number[]) {
+	return 1.0 - dot(embeddingA, embeddingB);
+}
+
+export async function findSimilarSentences(
+	embeddingModel: EmbeddingBackendModel,
+	query: string,
+	sentences: string[],
+	{ topK = 5 }: { topK: number }
+): Promise<number[]> {
+	const inputs = [query, ...sentences];
+
+	const embeddingEndpoint = await embeddingModel.getEndpoint();
+	const output = await embeddingEndpoint({ inputs })
+
+	const queryEmbedding: number[] = output[0];
+	const sentencesEmbeddings: number[][] = output.slice([1, inputs.length - 1]);
+
+	const distancesFromQuery: { distance: number; index: number }[] = [...sentencesEmbeddings].map(
+		(sentenceEmbedding: number[], index: number) => {
+			return {
+				distance: innerProduct(queryEmbedding, sentenceEmbedding),
+				index: index,
+			};
+		}
+	);
+
+	distancesFromQuery.sort((a, b) => {
+		return a.distance - b.distance;
+	});
+
+	// Return the indexes of the closest topK sentences
+	return distancesFromQuery.slice(0, topK).map((item) => item.index);
+}

src/lib/server/websearch/runWebSearch.ts

@@ -4,13 +4,11 @@ import type { WebSearch, WebSearchSource } from "$lib/types/WebSearch";
 import { generateQuery } from "$lib/server/websearch/generateQuery";
 import { parseWeb } from "$lib/server/websearch/parseWeb";
 import { chunk } from "$lib/utils/chunk";
-import {
-	MAX_SEQ_LEN as CHUNK_CAR_LEN,
-	findSimilarSentences,
-} from "$lib/server/websearch/sentenceSimilarity";
+import { findSimilarSentences } from "$lib/server/sentenceSimilarity";
 import type { Conversation } from "$lib/types/Conversation";
 import type { MessageUpdate } from "$lib/types/MessageUpdate";
 import { getWebSearchProvider } from "./searchWeb";
+import { embeddingModels } from "../embeddingModels";
 
 const MAX_N_PAGES_SCRAPE = 10 as const;
 const MAX_N_PAGES_EMBED = 5 as const;
@@ -57,6 +55,12 @@ export async function runWebSearch(
 			.filter(({ link }) => !DOMAIN_BLOCKLIST.some((el) => link.includes(el))) // filter out blocklist links
 			.slice(0, MAX_N_PAGES_SCRAPE); // limit to first 10 links only
 
+		// fetch the model
+		const embeddingModel = embeddingModels.find((m) => m.id === conv.embeddingModel);
+		if (!embeddingModel) {
+			throw new Error(`Embedding model ${conv.embeddingModel} not available anymore`);
+		}
+
 		let paragraphChunks: { source: WebSearchSource; text: string }[] = [];
 		if (webSearch.results.length > 0) {
 			appendUpdate("Browsing results");
@@ -72,7 +76,7 @@ export async function runWebSearch(
 					}
 				}
 				const MAX_N_CHUNKS = 100;
-				const texts = chunk(text, CHUNK_CAR_LEN).slice(0, MAX_N_CHUNKS);
+				const texts = chunk(text, embeddingModel.maxSequenceLength).slice(0, MAX_N_CHUNKS);
 				return texts.map((t) => ({ source: result, text: t }));
 			});
 			const nestedParagraphChunks = (await Promise.all(promises)).slice(0, MAX_N_PAGES_EMBED);
@@ -87,7 +91,7 @@ export async function runWebSearch(
 		appendUpdate("Extracting relevant information");
 		const topKClosestParagraphs = 8;
 		const texts = paragraphChunks.map(({ text }) => text);
-		const indices = await findSimilarSentences(prompt, texts, {
+		const indices = await findSimilarSentences(embeddingModel, prompt, texts, {
 			topK: topKClosestParagraphs,
 		});
 		webSearch.context = indices.map((idx) => texts[idx]).join("");