minor changes

huggingface · Feb 5, 2025 · 10436c7 · 10436c7
1 parent 4a7c7ad
commit 10436c7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/packages/hub/src/types/api/api-model.ts b/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,9 +18,8 @@ export interface ApiModelInfo {
 	downloadsAllTime: number;
 	files: string[];
 	gitalyUid: string;
-	inferenceProviderMapping: Record<
-		string,
-		{ providerId: string; status: "prod" | "staging"; task: PipelineType | "conversational" }
+	inferenceProviderMapping: Partial<
+		Record<string, { providerId: string; status: "prod" | "staging"; task: WidgetType }>
 	>;
 	lastAuthor: { email: string; user?: string };
 	lastModified: string; // convert to date

diff --git a/packages/inference/src/lib/makeRequestOptions.ts b/packages/inference/src/lib/makeRequestOptions.ts
@@ -172,9 +172,9 @@ async function mapModel(params: {
 	// 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 (params.provider in info.inferenceProviderMapping) {
-		const inferenceProviderMapping = info.inferenceProviderMapping[params.provider];
+	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}.`