Update conversational inference API snippets

packages/tasks/src/snippets/common.ts

@@ -0,0 +1,63 @@
+import type { ChatCompletionInputMessage, GenerationParameters } from "../tasks";
+
+export interface StringifyMessagesOptions {
+	sep: string;
+	start: string;
+	end: string;
+	attributeKeyQuotes?: boolean;
+	customContentEscaper?: (str: string) => string;
+}
+
+export function stringifyMessages(messages: ChatCompletionInputMessage[], opts: StringifyMessagesOptions): string {
+	const keyRole = opts.attributeKeyQuotes ? `"role"` : "role";
+	const keyContent = opts.attributeKeyQuotes ? `"role"` : "role";
+
+	const messagesStringified = messages.map(({ role, content }) => {
+		if (typeof content === "string") {
+			content = JSON.stringify(content).slice(1, -1);
+			if (opts.customContentEscaper) {
+				content = opts.customContentEscaper(content);
+			}
+			return `{ ${keyRole}: "${role}", ${keyContent}: "${content}" }`;
+		} else {
+			2;
+			content = content.map(({ image_url, text, type }) => ({
+				type,
+				image_url,
+				...(text ? { text: JSON.stringify(text).slice(1, -1) } : undefined),
+			}));
+			content = JSON.stringify(content).slice(1, -1);
+			if (opts.customContentEscaper) {
+				content = opts.customContentEscaper(content);
+			}
+			return `{ ${keyRole}: "${role}", ${keyContent}: ${content} }`;
+		}
+	});
+
+	return opts.start + messagesStringified.join(opts.sep) + opts.end;
+}
+
+type PartialGenerationParameters = Partial<Pick<GenerationParameters, "temperature" | "max_tokens" | "top_p">>;
+
+export interface StringifyGenerationConfigOptions {
+	sep: string;
+	start: string;
+	end: string;
+	attributeValueConnector: string;
+	attributeKeyQuotes?: boolean;
+}
+
+export function stringifyGenerationConfig(
+	config: PartialGenerationParameters,
+	opts: StringifyGenerationConfigOptions
+): string {
+	const quote = opts.attributeKeyQuotes ? `"` : "";
+
+	return (
+		opts.start +
+		Object.entries(config)
+			.map(([key, val]) => `${quote}${key}${quote}${opts.attributeValueConnector}${val}`)
+			.join(opts.sep) +
+		opts.end
+	);
+}

packages/tasks/src/snippets/curl.ts

@@ -1,36 +1,73 @@
 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 { ModelDataMinimal } from "./types.js";
+import type { InferenceSnippet, ModelDataMinimal } from "./types.js";
 
-export const snippetBasic = (model: ModelDataMinimal, accessToken: string): string =>
-	`curl https://api-inference.huggingface.co/models/${model.id} \\
+export const snippetBasic = (model: ModelDataMinimal, accessToken: string): InferenceSnippet => ({
+	content: `curl https://api-inference.huggingface.co/models/${model.id} \\
 	-X POST \\
 	-d '{"inputs": ${getModelInputSnippet(model, true)}}' \\
 	-H 'Content-Type: application/json' \\
-	-H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}"`;
+	-H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}"`,
+});
 
-export const snippetTextGeneration = (model: ModelDataMinimal, accessToken: string): string => {
+export const snippetTextGeneration = (
+	model: ModelDataMinimal,
+	accessToken: string,
+	opts?: {
+		streaming?: boolean;
+		messages?: ChatCompletionInputMessage[];
+		temperature?: GenerationParameters["temperature"];
+		max_tokens?: GenerationParameters["max_tokens"];
+		top_p?: GenerationParameters["top_p"];
+	}
+): InferenceSnippet => {
 	if (model.tags.includes("conversational")) {
 		// Conversational model detected, so we display a code snippet that features the Messages API
-		return `curl 'https://api-inference.huggingface.co/models/${model.id}/v1/chat/completions' \\
+		const streaming = opts?.streaming ?? true;
+		const messages: ChatCompletionInputMessage[] = opts?.messages ?? [
+			{ role: "user", content: "What is the capital of France?" },
+		];
+
+		const config = {
+			...(opts?.temperature ? { temperature: opts.temperature } : undefined),
+			max_tokens: opts?.max_tokens ?? 500,
+			...(opts?.top_p ? { top_p: opts.top_p } : undefined),
+		};
+		return {
+			content: `curl 'https://api-inference.huggingface.co/models/${model.id}/v1/chat/completions' \\
 -H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}" \\
 -H 'Content-Type: application/json' \\
--d '{
-	"model": "${model.id}",
-	"messages": [{"role": "user", "content": "What is the capital of France?"}],
-	"max_tokens": 500,
-	"stream": false
-}'
-`;
+--data '{
+    "model": "${model.id}",
+    "messages": ${stringifyMessages(messages, {
+			sep: ",\n\t\t",
+			start: `[\n\t\t`,
+			end: `\n\t]`,
+			attributeKeyQuotes: true,
+			customContentEscaper: (str) => str.replace(/'/g, "'\\''"),
+		})},
+    ${stringifyGenerationConfig(config, {
+			sep: ",\n    ",
+			start: "",
+			end: "",
+			attributeKeyQuotes: true,
+			attributeValueConnector: ": ",
+		})},
+    "stream": ${!!streaming}
+}'`,
+		};
 	} else {
 		return snippetBasic(model, accessToken);
 	}
 };
 
-export const snippetImageTextToTextGeneration = (model: ModelDataMinimal, accessToken: string): string => {
+export const snippetImageTextToTextGeneration = (model: ModelDataMinimal, accessToken: string): InferenceSnippet => {
 	if (model.tags.includes("conversational")) {
 		// Conversational model detected, so we display a code snippet that features the Messages API
-		return `curl 'https://api-inference.huggingface.co/models/${model.id}/v1/chat/completions' \\
+		return {
+			content: `curl 'https://api-inference.huggingface.co/models/${model.id}/v1/chat/completions' \\
 -H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}" \\
 -H 'Content-Type: application/json' \\
 -d '{
@@ -47,26 +84,34 @@ export const snippetImageTextToTextGeneration = (model: ModelDataMinimal, access
 	"max_tokens": 500,
 	"stream": false
 }'
-`;
+`,
+		};
 	} else {
 		return snippetBasic(model, accessToken);
 	}
 };
 
-export const snippetZeroShotClassification = (model: ModelDataMinimal, accessToken: string): string =>
-	`curl https://api-inference.huggingface.co/models/${model.id} \\
+export const snippetZeroShotClassification = (model: ModelDataMinimal, accessToken: string): InferenceSnippet => ({
+	content: `curl https://api-inference.huggingface.co/models/${model.id} \\
 	-X POST \\
 	-d '{"inputs": ${getModelInputSnippet(model, true)}, "parameters": {"candidate_labels": ["refund", "legal", "faq"]}}' \\
 	-H 'Content-Type: application/json' \\
-	-H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}"`;
+	-H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}"`,
+});
 
-export const snippetFile = (model: ModelDataMinimal, accessToken: string): string =>
-	`curl https://api-inference.huggingface.co/models/${model.id} \\
+export const snippetFile = (model: ModelDataMinimal, accessToken: string): InferenceSnippet => ({
+	content: `curl https://api-inference.huggingface.co/models/${model.id} \\
 	-X POST \\
 	--data-binary '@${getModelInputSnippet(model, true, true)}' \\
-	-H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}"`;
+	-H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}"`,
+});
 
-export const curlSnippets: Partial<Record<PipelineType, (model: ModelDataMinimal, accessToken: string) => string>> = {
+export const curlSnippets: Partial<
+	Record<
+		PipelineType,
+		(model: ModelDataMinimal, accessToken: string, opts?: Record<string, unknown>) => InferenceSnippet
+	>
+> = {
 	// Same order as in js/src/lib/interfaces/Types.ts
 	"text-classification": snippetBasic,
 	"token-classification": snippetBasic,
@@ -93,10 +138,10 @@ export const curlSnippets: Partial<Record<PipelineType, (model: ModelDataMinimal
 	"image-segmentation": snippetFile,
 };
 
-export function getCurlInferenceSnippet(model: ModelDataMinimal, accessToken: string): string {
+export function getCurlInferenceSnippet(model: ModelDataMinimal, accessToken: string): InferenceSnippet {
 	return model.pipeline_tag && model.pipeline_tag in curlSnippets
-		? curlSnippets[model.pipeline_tag]?.(model, accessToken) ?? ""
-		: "";
+		? curlSnippets[model.pipeline_tag]?.(model, accessToken) ?? { content: "" }
+		: { content: "" };
 }
 
 export function hasCurlInferenceSnippet(model: Pick<ModelDataMinimal, "pipeline_tag">): boolean {