(WIP) Dynamic inference provider mapping

packages/hub/src/lib/list-models.ts

@@ -25,6 +25,7 @@ export const MODEL_EXPANDABLE_KEYS = [
 	"downloadsAllTime",
 	"gated",
 	"gitalyUid",
+	"inferenceProviderMapping",
 	"lastModified",
 	"library_name",
 	"likes",

packages/hub/src/types/api/api-model.ts

@@ -1,4 +1,4 @@
-import type { ModelLibraryKey, TransformersInfo } from "@huggingface/tasks";
+import type { ModelLibraryKey, TransformersInfo, WidgetType } from "@huggingface/tasks";
 import type { License, PipelineType } from "../public";
 
 export interface ApiModelInfo {
@@ -18,6 +18,9 @@ export interface ApiModelInfo {
 	downloadsAllTime: number;
 	files: string[];
 	gitalyUid: string;
+	inferenceProviderMapping: Partial<
+		Record<string, { providerId: string; status: "prod" | "staging"; task: WidgetType }>
+	>;
 	lastAuthor: { email: string; user?: string };
 	lastModified: string; // convert to date
 	library_name?: ModelLibraryKey;

packages/inference/package.json

@@ -52,6 +52,7 @@
 		"check": "tsc"
 	},
 	"dependencies": {
+		"@huggingface/hub": "workspace:^",
 		"@huggingface/tasks": "workspace:^"
 	},
 	"devDependencies": {

packages/inference/src/lib/makeRequestOptions.ts

@@ -1,3 +1,4 @@
+import { modelInfo } from "@huggingface/hub";
 import type { WidgetType } from "@huggingface/tasks";
 import { HF_HUB_URL } from "../config";
 import { FAL_AI_API_BASE_URL, FAL_AI_SUPPORTED_MODEL_IDS } from "../providers/fal-ai";
@@ -53,13 +54,13 @@ export async function makeRequestOptions(
 	let model: string;
 	if (!maybeModel) {
 		if (taskHint) {
-			model = mapModel({ model: await loadDefaultModel(taskHint), provider, taskHint, chatCompletion });
+			model = await mapModel({ model: await loadDefaultModel(taskHint), provider, taskHint, chatCompletion });
 		} else {
 			throw new Error("No model provided, and no default model found for this task");
 			/// TODO : change error message ^
 		}
 	} else {
-		model = mapModel({ model: maybeModel, provider, taskHint, chatCompletion });
+		model = await mapModel({ model: maybeModel, provider, taskHint, chatCompletion });
 	}
 
 	/// If accessToken is passed, it should take precedence over includeCredentials
@@ -153,12 +154,12 @@ export async function makeRequestOptions(
 	return { url, info };
 }
 
-function mapModel(params: {
+async function mapModel(params: {
 	model: string;
 	provider: InferenceProvider;
 	taskHint: InferenceTask | undefined;
 	chatCompletion: boolean | undefined;
-}): string {
+}): Promise<string> {
 	if (params.provider === "hf-inference") {
 		return params.model;
 	}
@@ -167,7 +168,29 @@ function mapModel(params: {
 	}
 	const task: WidgetType =
 		params.taskHint === "text-generation" && params.chatCompletion ? "conversational" : params.taskHint;
-	const model = (() => {
+
+	// TODO: cache this call
+	const info = await modelInfo({ name: params.model, additionalFields: ["inferenceProviderMapping"] });
+
+	const inferenceProviderMapping = info.inferenceProviderMapping[params.provider];
+	// If provider listed => takes precedence over hard-coded mapping
+	if (inferenceProviderMapping) {
+		if (inferenceProviderMapping.task !== task) {
+			throw new Error(
+				`Model ${params.model} is not supported for task ${task} and provider ${params.provider}. Supported task: ${inferenceProviderMapping.task}.`
+			);
+		}
+		if (inferenceProviderMapping.status === "staging") {
+			console.warn(
+				`Model ${params.model} is in staging for provider ${params.provider}. Use it only for test purposes.`
+			);
+		}
+		// TODO: how is it handled server-side if model has multiple tasks (e.g. `text-generation` + `conversational`)?
+		return inferenceProviderMapping.providerId;
+	}
+
+	// Otherwise, default to hard-coded mapping
+	const modelFromMapping = (() => {
 		switch (params.provider) {
 			case "fal-ai":
 				return FAL_AI_SUPPORTED_MODEL_IDS[task]?.[params.model];
@@ -179,11 +202,10 @@ function mapModel(params: {
 				return TOGETHER_SUPPORTED_MODEL_IDS[task]?.[params.model];
 		}
 	})();
-
-	if (!model) {
-		throw new Error(`Model ${params.model} is not supported for task ${task} and provider ${params.provider}`);
+	if (!modelFromMapping) {
+		throw new Error(`Model ${params.model} is not supported for task ${task} and provider ${params.provider}.`);
 	}
-	return model;
+	return modelFromMapping;
 }
 
 function makeUrl(params: {

packages/inference/src/providers/sambanova.ts

@@ -7,7 +7,7 @@ type SambanovaId = string;
 export const SAMBANOVA_SUPPORTED_MODEL_IDS: ProviderMapping<SambanovaId> = {
 	/** Chat completion / conversational */
 	conversational: {
-		"allenai/Llama-3.1-Tulu-3-405B":"Llama-3.1-Tulu-3-405B",
+		"allenai/Llama-3.1-Tulu-3-405B": "Llama-3.1-Tulu-3-405B",
 		"deepseek-ai/DeepSeek-R1-Distill-Llama-70B": "DeepSeek-R1-Distill-Llama-70B",
 		"Qwen/Qwen2.5-Coder-32B-Instruct": "Qwen2.5-Coder-32B-Instruct",
 		"Qwen/Qwen2.5-72B-Instruct": "Qwen2.5-72B-Instruct",