@huggingface/tasks 0.14.0 → 0.15.0

package/src/snippets/python.ts

@@ -1,17 +1,23 @@
+import { HF_HUB_INFERENCE_PROXY_TEMPLATE, openAIbaseUrl, type InferenceProvider } from "../inference-providers.js";
 import type { PipelineType } from "../pipelines.js";
 import type { ChatCompletionInputMessage, GenerationParameters } from "../tasks/index.js";
 import { stringifyGenerationConfig, stringifyMessages } from "./common.js";
 import { getModelInputSnippet } from "./inputs.js";
 import type { InferenceSnippet, ModelDataMinimal } from "./types.js";
 
-const snippetImportInferenceClient = (model: ModelDataMinimal, accessToken: string): string =>
-	`from huggingface_hub import InferenceClient
-client = InferenceClient("${model.id}", token="${accessToken || "{API_TOKEN}"}")
-`;
+const snippetImportInferenceClient = (accessToken: string, provider: InferenceProvider): string =>
+	`\
+from huggingface_hub import InferenceClient
+
+client = InferenceClient(
+	provider="${provider}",
+	api_key="${accessToken || "{API_TOKEN}"}"
+)`;
 
 export const snippetConversational = (
 	model: ModelDataMinimal,
 	accessToken: string,
+	provider: InferenceProvider,
 	opts?: {
 		streaming?: boolean;
 		messages?: ChatCompletionInputMessage[];
@@ -39,9 +45,8 @@ export const snippetConversational = (
 		return [
 			{
 				client: "huggingface_hub",
-				content: `from huggingface_hub import InferenceClient
-
-client = InferenceClient(api_key="${accessToken || "{API_TOKEN}"}")
+				content: `\
+${snippetImportInferenceClient(accessToken, provider)}
 
 messages = ${messagesStr}
 
@@ -60,7 +65,7 @@ for chunk in stream:
 				content: `from openai import OpenAI
 
 client = OpenAI(
-	base_url="https://api-inference.huggingface.co/v1/",
+	base_url="${openAIbaseUrl(provider)}",
 	api_key="${accessToken || "{API_TOKEN}"}"
 )
 
@@ -81,9 +86,8 @@ for chunk in stream:
 		return [
 			{
 				client: "huggingface_hub",
-				content: `from huggingface_hub import InferenceClient
-
-client = InferenceClient(api_key="${accessToken || "{API_TOKEN}"}")
+				content: `\
+${snippetImportInferenceClient(accessToken, provider)}
 
 messages = ${messagesStr}
 
@@ -100,7 +104,7 @@ print(completion.choices[0].message)`,
 				content: `from openai import OpenAI
 
 client = OpenAI(
-	base_url="https://api-inference.huggingface.co/v1/",
+	base_url="${openAIbaseUrl(provider)}",
 	api_key="${accessToken || "{API_TOKEN}"}"
 )
 
@@ -118,8 +122,11 @@ print(completion.choices[0].message)`,
 	}
 };
 
-export const snippetZeroShotClassification = (model: ModelDataMinimal): InferenceSnippet => ({
-	content: `def query(payload):
+export const snippetZeroShotClassification = (model: ModelDataMinimal): InferenceSnippet[] => {
+	return [
+		{
+			client: "requests",
+			content: `def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
 	return response.json()
 
@@ -127,10 +134,15 @@ output = query({
     "inputs": ${getModelInputSnippet(model)},
     "parameters": {"candidate_labels": ["refund", "legal", "faq"]},
 })`,
-});
+		},
+	];
+};
 
-export const snippetZeroShotImageClassification = (model: ModelDataMinimal): InferenceSnippet => ({
-	content: `def query(data):
+export const snippetZeroShotImageClassification = (model: ModelDataMinimal): InferenceSnippet[] => {
+	return [
+		{
+			client: "requests",
+			content: `def query(data):
 	with open(data["image_path"], "rb") as f:
 		img = f.read()
 	payload={
@@ -144,40 +156,85 @@ output = query({
     "image_path": ${getModelInputSnippet(model)},
     "parameters": {"candidate_labels": ["cat", "dog", "llama"]},
 })`,
-});
+		},
+	];
+};
 
-export const snippetBasic = (model: ModelDataMinimal): InferenceSnippet => ({
-	content: `def query(payload):
+export const snippetBasic = (model: ModelDataMinimal): InferenceSnippet[] => {
+	return [
+		{
+			client: "requests",
+			content: `def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
 	return response.json()
 	
 output = query({
 	"inputs": ${getModelInputSnippet(model)},
 })`,
-});
-
-export const snippetFile = (model: ModelDataMinimal): InferenceSnippet => ({
-	content: `def query(filename):
-    with open(filename, "rb") as f:
-        data = f.read()
-    response = requests.post(API_URL, headers=headers, data=data)
-    return response.json()
-
-output = query(${getModelInputSnippet(model)})`,
-});
-
-export const snippetTextToImage = (model: ModelDataMinimal, accessToken: string): InferenceSnippet[] => [
-	{
-		client: "huggingface_hub",
-		content: `${snippetImportInferenceClient(model, accessToken)}
+		},
+	];
+};
+
+export const snippetFile = (model: ModelDataMinimal): InferenceSnippet[] => {
+	return [
+		{
+			client: "requests",
+			content: `def query(filename):
+		with open(filename, "rb") as f:
+			data = f.read()
+		response = requests.post(API_URL, headers=headers, data=data)
+		return response.json()
+	
+	output = query(${getModelInputSnippet(model)})`,
+		},
+	];
+};
+
+export const snippetTextToImage = (
+	model: ModelDataMinimal,
+	accessToken: string,
+	provider: InferenceProvider
+): InferenceSnippet[] => {
+	return [
+		{
+			client: "huggingface_hub",
+			content: `\
+${snippetImportInferenceClient(accessToken, provider)}
+
 # output is a PIL.Image object
-image = client.text_to_image(${getModelInputSnippet(model)})`,
+image = client.text_to_image(
+	${getModelInputSnippet(model)},
+	model="${model.id}"
+)`,
+		},
+		...(provider === "fal-ai"
+			? [
+					{
+						client: "fal-client",
+						content: `\
+import fal_client
+
+result = fal_client.subscribe(
+	# replace with correct id from fal.ai
+	"fal-ai/${model.id}",
+	arguments={
+		"prompt": ${getModelInputSnippet(model)},
 	},
-	{
-		client: "requests",
-		content: `def query(payload):
+)
+print(result)
+`,
+					},
+			  ]
+			: []),
+		...(provider === "hf-inference"
+			? [
+					{
+						client: "requests",
+						content: `\
+def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
 	return response.content
+
 image_bytes = query({
 	"inputs": ${getModelInputSnippet(model)},
 })
@@ -186,25 +243,35 @@ image_bytes = query({
 import io
 from PIL import Image
 image = Image.open(io.BytesIO(image_bytes))`,
-	},
-];
+					},
+			  ]
+			: []),
+	];
+};
 
-export const snippetTabular = (model: ModelDataMinimal): InferenceSnippet => ({
-	content: `def query(payload):
-	response = requests.post(API_URL, headers=headers, json=payload)
-	return response.content
-response = query({
-	"inputs": {"data": ${getModelInputSnippet(model)}},
-})`,
-});
+export const snippetTabular = (model: ModelDataMinimal): InferenceSnippet[] => {
+	return [
+		{
+			client: "requests",
+			content: `def query(payload):
+		response = requests.post(API_URL, headers=headers, json=payload)
+		return response.content
+	response = query({
+		"inputs": {"data": ${getModelInputSnippet(model)}},
+	})`,
+		},
+	];
+};
 
-export const snippetTextToAudio = (model: ModelDataMinimal): InferenceSnippet => {
+export const snippetTextToAudio = (model: ModelDataMinimal): InferenceSnippet[] => {
 	// Transformers TTS pipeline and api-inference-community (AIC) pipeline outputs are diverged
 	// with the latest update to inference-api (IA).
 	// Transformers IA returns a byte object (wav file), whereas AIC returns wav and sampling_rate.
 	if (model.library_name === "transformers") {
-		return {
-			content: `def query(payload):
+		return [
+			{
+				client: "requests",
+				content: `def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
 	return response.content
 
@@ -214,10 +281,13 @@ audio_bytes = query({
 # You can access the audio with IPython.display for example
 from IPython.display import Audio
 Audio(audio_bytes)`,
-		};
+			},
+		];
 	} else {
-		return {
-			content: `def query(payload):
+		return [
+			{
+				client: "requests",
+				content: `def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
 	return response.json()
 	
@@ -227,12 +297,16 @@ audio, sampling_rate = query({
 # You can access the audio with IPython.display for example
 from IPython.display import Audio
 Audio(audio, rate=sampling_rate)`,
-		};
+			},
+		];
 	}
 };
 
-export const snippetDocumentQuestionAnswering = (model: ModelDataMinimal): InferenceSnippet => ({
-	content: `def query(payload):
+export const snippetDocumentQuestionAnswering = (model: ModelDataMinimal): InferenceSnippet[] => {
+	return [
+		{
+			client: "requests",
+			content: `def query(payload):
  	with open(payload["image"], "rb") as f:
   		img = f.read()
 		payload["image"] = base64.b64encode(img).decode("utf-8")  
@@ -242,7 +316,9 @@ export const snippetDocumentQuestionAnswering = (model: ModelDataMinimal): Infer
 output = query({
     "inputs": ${getModelInputSnippet(model)},
 })`,
-});
+		},
+	];
+};
 
 export const pythonSnippets: Partial<
 	Record<
@@ -250,8 +326,9 @@ export const pythonSnippets: Partial<
 		(
 			model: ModelDataMinimal,
 			accessToken: string,
+			provider: InferenceProvider,
 			opts?: Record<string, unknown>
-		) => InferenceSnippet | InferenceSnippet[]
+		) => InferenceSnippet[]
 	>
 > = {
 	// Same order as in tasks/src/pipelines.ts
@@ -287,35 +364,37 @@ export const pythonSnippets: Partial<
 export function getPythonInferenceSnippet(
 	model: ModelDataMinimal,
 	accessToken: string,
+	provider: InferenceProvider,
 	opts?: Record<string, unknown>
-): InferenceSnippet | InferenceSnippet[] {
+): InferenceSnippet[] {
 	if (model.tags.includes("conversational")) {
 		// Conversational model detected, so we display a code snippet that features the Messages API
-		return snippetConversational(model, accessToken, opts);
+		return snippetConversational(model, accessToken, provider, opts);
 	} else {
-		let snippets =
+		const snippets =
 			model.pipeline_tag && model.pipeline_tag in pythonSnippets
-				? pythonSnippets[model.pipeline_tag]?.(model, accessToken) ?? { content: "" }
-				: { content: "" };
+				? pythonSnippets[model.pipeline_tag]?.(model, accessToken, provider) ?? []
+				: [];
 
-		snippets = Array.isArray(snippets) ? snippets : [snippets];
+		const baseUrl =
+			provider === "hf-inference"
+				? `https://api-inference.huggingface.co/models/${model.id}`
+				: HF_HUB_INFERENCE_PROXY_TEMPLATE.replace("{{PROVIDER}}", provider);
 
 		return snippets.map((snippet) => {
 			return {
 				...snippet,
-				content: snippet.content.includes("requests")
-					? `import requests
+				content:
+					snippet.client === "requests"
+						? `\
+import requests
 
-API_URL = "https://api-inference.huggingface.co/models/${model.id}"
+API_URL = "${baseUrl}"
 headers = {"Authorization": ${accessToken ? `"Bearer ${accessToken}"` : `f"Bearer {API_TOKEN}"`}}
 
 ${snippet.content}`
-					: snippet.content,
+						: snippet.content,
 			};
 		});
 	}
 }
-
-export function hasPythonInferenceSnippet(model: ModelDataMinimal): boolean {
-	return !!model.pipeline_tag && model.pipeline_tag in pythonSnippets;
-}

package/src/snippets/types.ts

@@ -12,5 +12,5 @@ export type ModelDataMinimal = Pick<
 
 export interface InferenceSnippet {
 	content: string;
-	client?: string; // for instance: `client` could be `huggingface_hub` or `openai` client for Python snippets
+	client: string; // for instance: `client` could be `huggingface_hub` or `openai` client for Python snippets
 }