Generate feature-extraction specs from TEI

packages/tasks/package.json

@@ -28,7 +28,8 @@
 		"prepare": "pnpm run build",
 		"check": "tsc",
 		"inference-codegen": "tsx scripts/inference-codegen.ts && prettier --write src/tasks/*/inference.ts",
-		"inference-tgi-import": "tsx scripts/inference-tgi-import.ts && prettier --write src/tasks/text-generation/spec/*.json && prettier --write src/tasks/chat-completion/spec/*.json"
+		"inference-tgi-import": "tsx scripts/inference-tgi-import.ts && prettier --write src/tasks/text-generation/spec/*.json && prettier --write src/tasks/chat-completion/spec/*.json",
+		"inference-tei-import": "tsx scripts/inference-tei-import.ts && prettier --write src/tasks/feature-extraction/spec/*.json"
 	},
 	"type": "module",
 	"files": [

packages/tasks/scripts/inference-tei-import.ts

@@ -0,0 +1,111 @@
+/*
+ * Fetches TGI specs and generated JSON schema for input, output and stream_output of
+ * text-generation and chat-completion tasks.
+ * See https://huggingface.github.io/text-generation-inference/
+ */
+import fs from "fs/promises";
+import fetch from "node-fetch";
+import * as path from "node:path/posix";
+import { existsSync as pathExists } from "node:fs";
+import type { JsonObject, JsonValue } from "type-fest";
+
+const URL = "https://huggingface.github.io/text-embeddings-inference/openapi.json";
+
+const rootDirFinder = function (): string {
+	let currentPath = path.normalize(import.meta.url);
+
+	while (currentPath !== "/") {
+		if (pathExists(path.join(currentPath, "package.json"))) {
+			return currentPath;
+		}
+
+		currentPath = path.normalize(path.join(currentPath, ".."));
+	}
+
+	return "/";
+};
+
+const rootDir = rootDirFinder();
+const tasksDir = path.join(rootDir, "src", "tasks");
+
+function toCamelCase(str: string, joiner = "") {
+	return str
+		.split(/[-_]/)
+		.map((part) => part.charAt(0).toUpperCase() + part.slice(1))
+		.join(joiner);
+}
+
+async function _extractAndAdapt(task: string, mainComponentName: string, type: "input" | "output") {
+	console.debug(`✨ Importing`, task, type);
+
+	console.debug("   📥 Fetching TEI specs");
+	const response = await fetch(URL);
+	// eslint-disable-next-line @typescript-eslint/no-explicit-any
+	const openapi = (await response.json()) as any;
+	// eslint-disable-next-line @typescript-eslint/no-explicit-any
+	const components: Record<string, any> = openapi["components"]["schemas"];
+
+	// e.g. TextGeneration
+	const camelName = toCamelCase(task);
+	// e.g. TextGenerationInput
+	const camelFullName = camelName + toCamelCase(type);
+	const mainComponent = components[mainComponentName];
+	const filteredComponents: Record<string, JsonObject> = {};
+
+	function _scan(data: JsonValue) {
+		if (Array.isArray(data) || data instanceof Array) {
+			for (const item of data) {
+				_scan(item);
+			}
+		} else if (data && typeof data === "object") {
+			for (const key of Object.keys(data)) {
+				if (key === "$ref" && typeof data[key] === "string") {
+					// Verify reference exists
+					const ref = (data[key] as string).split("/").pop() ?? "";
+					if (!components[ref]) {
+						throw new Error(`Reference not found in components: ${data[key]}`);
+					}
+
+					// Add reference to components to export (and scan it too)
+					const newRef = camelFullName + ref.replace(camelName, "");
+					if (!filteredComponents[newRef]) {
+						components[ref]["title"] = newRef; // Rename title to avoid conflicts
+						filteredComponents[newRef] = components[ref];
+						_scan(components[ref]);
+					}
+
+					// Updating the reference to new format
+					data[key] = `#/$defs/${newRef}`;
+				} else {
+					_scan(data[key]);
+				}
+			}
+		}
+	}
+
+	console.debug("   📦 Packaging jsonschema");
+	_scan(mainComponent);
+
+	const prettyName = toCamelCase(task, " ") + " " + toCamelCase(type, " ");
+	const inputSchema = {
+		$id: `/inference/schemas/${task}/${type}.json`,
+		$schema: "http://json-schema.org/draft-06/schema#",
+		description:
+			prettyName +
+			".\n\nAuto-generated from TEI specs." +
+			"\nFor more details, check out https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts.",
+		title: camelFullName,
+		type: "object",
+		required: mainComponent["required"],
+		properties: mainComponent["properties"],
+		$defs: filteredComponents,
+	};
+
+	const specPath = path.join(tasksDir, task, "spec", `${type}.json`);
+	console.debug("   📂 Exporting", specPath);
+	await fs.writeFile(specPath, JSON.stringify(inputSchema, null, 4));
+}
+
+await _extractAndAdapt("feature-extraction", "EmbedRequest", "input");
+await _extractAndAdapt("feature-extraction", "EmbedResponse", "output");
+console.debug("✅ All done!");

packages/tasks/src/tasks/feature-extraction/inference.ts

@@ -4,19 +4,20 @@
  * Using src/scripts/inference-codegen
  */
 
-export type FeatureExtractionOutput = unknown[];
+export interface FeatureExtractionOutput { [key: string]: unknown }
 
 /**
- * Inputs for Text Embedding inference
+ * Feature Extraction Input.
+ *
+ * Auto-generated from TEI specs.
+ * For more details, check out
+ * https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts.
  */
 export interface FeatureExtractionInput {
-	/**
-	 * The text to get the embeddings of
-	 */
-	inputs: string;
-	/**
-	 * Additional inference parameters
-	 */
-	parameters?: { [key: string]: unknown };
+	inputs: FeatureExtractionInputInput;
+	normalize?: boolean;
+	truncate?: boolean;
 	[property: string]: unknown;
 }
+
+export type FeatureExtractionInputInput = string[] | string;

packages/tasks/src/tasks/feature-extraction/spec/input.json

@@ -1,26 +1,39 @@
 {
 	"$id": "/inference/schemas/feature-extraction/input.json",
 	"$schema": "http://json-schema.org/draft-06/schema#",
-	"description": "Inputs for Text Embedding inference",
+	"description": "Feature Extraction Input.\n\nAuto-generated from TEI specs.\nFor more details, check out https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts.",
 	"title": "FeatureExtractionInput",
 	"type": "object",
+	"required": ["inputs"],
 	"properties": {
 		"inputs": {
-			"description": "The text to get the embeddings of",
-			"type": "string"
+			"$ref": "#/$defs/FeatureExtractionInputInput"
 		},
-		"parameters": {
-			"description": "Additional inference parameters",
-			"$ref": "#/$defs/FeatureExtractionParameters"
+		"normalize": {
+			"type": "boolean",
+			"default": "true",
+			"example": "true"
+		},
+		"truncate": {
+			"type": "boolean",
+			"default": "false",
+			"example": "false"
 		}
 	},
 	"$defs": {
-		"FeatureExtractionParameters": {
-			"title": "FeatureExtractionParameters",
-			"description": "Additional inference parameters for Feature Extraction",
-			"type": "object",
-			"properties": {}
+		"FeatureExtractionInputInput": {
+			"oneOf": [
+				{
+					"type": "string"
+				},
+				{
+					"type": "array",
+					"items": {
+						"type": "string"
+					}
+				}
+			],
+			"title": "FeatureExtractionInputInput"
 		}
-	},
-	"required": ["inputs"]
+	}
 }

packages/tasks/src/tasks/feature-extraction/spec/output.json

@@ -1,7 +1,8 @@
 {
 	"$id": "/inference/schemas/feature-extraction/output.json",
 	"$schema": "http://json-schema.org/draft-06/schema#",
-	"description": "The embedding for the input text, as a nested list (tensor) of floats",
-	"type": "array",
-	"title": "FeatureExtractionOutput"
+	"description": "Feature Extraction Output.\n\nAuto-generated from TEI specs.\nFor more details, check out https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts.",
+	"title": "FeatureExtractionOutput",
+	"type": "object",
+	"$defs": {}
 }