Harmonize inferenceProviderMapping additional parameter in modelInfo / listModel

Author: Wauplin

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

@@ -115,4 +115,23 @@ describe("listModels", () => {
 
 		expect(count).to.equal(10);
 	});
+
+	it("should list deepseek-ai models with inference provider mapping", async () => {
+		let count = 0;
+		for await (const entry of listModels({
+			search: { owner: "deepseek-ai" },
+			additionalFields: ["inferenceProviderMapping"],
+			limit: 1,
+		})) {
+			count++;
+			expect(entry.inferenceProviderMapping).to.be.an("array").that.is.not.empty;
+			for (const item of entry.inferenceProviderMapping ?? []) {
+				expect(item).to.have.property("provider").that.is.a("string").and.is.not.empty;
+				expect(item).to.have.property("hfModelId").that.is.a("string").and.is.not.empty;
+				expect(item).to.have.property("providerId").that.is.a("string").and.is.not.empty;
+			}
+		}
+
+		expect(count).to.equal(1);
+	});
 });

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

@@ -5,6 +5,7 @@ import type { CredentialsParams, PipelineType } from "../types/public";
 import { checkCredentials } from "../utils/checkCredentials";
 import { parseLinkHeader } from "../utils/parseLinkHeader";
 import { pick } from "../utils/pick";
+import { normalizeInferenceProviderMapping } from "../utils/normalizeInferenceProviderMapping";
 
 export const MODEL_EXPAND_KEYS = [
 	"pipeline_tag",
@@ -113,8 +114,20 @@ export async function* listModels<
 		const items: ApiModelInfo[] = await res.json();
 
 		for (const item of items) {
+			// Handle inferenceProviderMapping normalization
+			const normalizedItem = { ...item };
+			if (
+				(params?.additionalFields as string[])?.includes("inferenceProviderMapping") &&
+				item.inferenceProviderMapping
+			) {
+				normalizedItem.inferenceProviderMapping = normalizeInferenceProviderMapping(
+					item.id,
+					item.inferenceProviderMapping
+				);
+			}
+
 			yield {
-				...(params?.additionalFields && pick(item, params.additionalFields)),
+				...(params?.additionalFields && pick(normalizedItem, params.additionalFields)),
 				id: item._id,
 				name: item.id,
 				private: item.private,

packages/hub/src/lib/model-info.spec.ts

@@ -56,4 +56,20 @@ describe("modelInfo", () => {
 			sha: "f27b190eeac4c2302d24068eabf5e9d6044389ae",
 		});
 	});
+
+	it("should return model info deepseek-ai models with inference provider mapping", async () => {
+		const info = await modelInfo({
+			name: "deepseek-ai/DeepSeek-R1-0528",
+			additionalFields: ["inferenceProviderMapping"],
+		});
+
+		expect(info.inferenceProviderMapping).toBeDefined();
+		expect(info.inferenceProviderMapping).toBeInstanceOf(Array);
+		expect(info.inferenceProviderMapping?.length).toBeGreaterThan(0);
+		info.inferenceProviderMapping?.forEach((item) => {
+			expect(item).toHaveProperty("provider");
+			expect(item).toHaveProperty("hf_model_id", "deepseek-ai/DeepSeek-R1-0528");
+			expect(item).toHaveProperty("provider_id");
+		});
+	});
 });

packages/hub/src/lib/model-info.ts

@@ -4,6 +4,7 @@ import type { ApiModelInfo } from "../types/api/api-model";
 import type { CredentialsParams } from "../types/public";
 import { checkCredentials } from "../utils/checkCredentials";
 import { pick } from "../utils/pick";
+import { normalizeInferenceProviderMapping } from "../utils/normalizeInferenceProviderMapping";
 import { MODEL_EXPAND_KEYS, type MODEL_EXPANDABLE_KEYS, type ModelEntry } from "./list-models";
 
 export async function modelInfo<
@@ -48,8 +49,14 @@ export async function modelInfo<
 
 	const data = await response.json();
 
+	// Handle inferenceProviderMapping normalization
+	const normalizedData = { ...data };
+	if ((params?.additionalFields as string[])?.includes("inferenceProviderMapping") && data.inferenceProviderMapping) {
+		normalizedData.inferenceProviderMapping = normalizeInferenceProviderMapping(data.id, data.inferenceProviderMapping);
+	}
+
 	return {
-		...(params?.additionalFields && pick(data, params.additionalFields)),
+		...(params?.additionalFields && pick(normalizedData, params.additionalFields)),
 		id: data._id,
 		name: data.id,
 		private: data.private,

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

@@ -18,9 +18,7 @@ export interface ApiModelInfo {
 	downloadsAllTime: number;
 	files: string[];
 	gitalyUid: string;
-	inferenceProviderMapping: Partial<
-		Record<string, { providerId: string; status: "live" | "staging"; task: WidgetType }>
-	>;
+	inferenceProviderMapping?: ApiModelInferenceProviderMappingEntry[];
 	lastAuthor: { email: string; user?: string };
 	lastModified: string; // convert to date
 	library_name?: ModelLibraryKey;
@@ -271,3 +269,14 @@ export interface ApiModelMetadata {
 	extra_gated_description?: string;
 	extra_gated_button_content?: string;
 }
+
+export interface ApiModelInferenceProviderMappingEntry {
+	provider: string; // Provider name
+	hf_model_id: string; // ID of the model on the Hugging Face Hub
+	provider_id: string; // ID of the model on the provider's side
+	status: "live" | "staging";
+	task: WidgetType;
+	adapter?: string;
+	adapter_weights_path?: string;
+	type?: "single-file" | "tag-filter";
+}

packages/hub/src/utils/normalizeInferenceProviderMapping.ts

@@ -0,0 +1,36 @@
+import type { WidgetType } from "@huggingface/tasks";
+import type { ApiModelInferenceProviderMappingEntry } from "../types/api/api-model";
+
+/**
+ * Normalize inferenceProviderMapping to always return an array format.
+ *
+ * Little hack to simplify Inference Providers logic and make it backward and forward compatible.
+ * Right now, API returns a dict on model-info and a list on list-models. Let's harmonize to list.
+ */
+export function normalizeInferenceProviderMapping(
+	hf_model_id: string,
+	inferenceProviderMapping?:
+		| ApiModelInferenceProviderMappingEntry[]
+		| Record<string, { providerId: string; status: "live" | "staging"; task: WidgetType }>
+): ApiModelInferenceProviderMappingEntry[] {
+	if (!inferenceProviderMapping) {
+		return [];
+	}
+
+	// If it's already an array, return it as is
+	if (Array.isArray(inferenceProviderMapping)) {
+		return inferenceProviderMapping.map((entry) => ({
+			...entry,
+			hf_model_id,
+		}));
+	}
+
+	// Convert mapping to array format
+	return Object.entries(inferenceProviderMapping).map(([provider, mapping]) => ({
+		provider,
+		hf_model_id,
+		provider_id: mapping.providerId,
+		status: mapping.status,
+		task: mapping.task,
+	}));
+}