curl

Author: mishig25

packages/tasks/src/snippets/curl.ts

@@ -1,6 +1,12 @@
 import type { PipelineType } from "../pipelines.js";
+import type { ChatCompletionInputMessage, GenerationParameters } from "../tasks/index.js";
 import { getModelInputSnippet } from "./inputs.js";
-import type { InferenceSnippet, ModelDataMinimal } from "./types.js";
+import type {
+	GenerationConfigFormatter,
+	GenerationMessagesFormatter,
+	InferenceSnippet,
+	ModelDataMinimal,
+} from "./types.js";
 
 export const snippetBasic = (model: ModelDataMinimal, accessToken: string): InferenceSnippet => ({
 	content: `curl https://api-inference.huggingface.co/models/${model.id} \\
@@ -10,20 +16,58 @@ export const snippetBasic = (model: ModelDataMinimal, accessToken: string): Infe
 	-H "Authorization: Bearer ${accessToken || `{API_TOKEN}`}"`,
 });
 
-export const snippetTextGeneration = (model: ModelDataMinimal, accessToken: string): InferenceSnippet => {
+const formatGenerationMessages: GenerationMessagesFormatter = ({ messages, sep, start, end }) =>
+	start +
+	messages
+		.map(({ role, content }) => {
+			// escape single quotes since single quotes is used to define http post body inside curl requests
+			// TODO: handle the case below
+			content = content?.replace(/'/g, "'\\''");
+			return `{ "role": "${role}", "content": "${content}" }`;
+		})
+		.join(sep) +
+	end;
+
+const formatGenerationConfig: GenerationConfigFormatter = ({ config, sep, start, end }) =>
+	start +
+	Object.entries(config)
+		.map(([key, val]) => `"${key}": ${val}`)
+		.join(sep) +
+	end;
+
+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
+		const streaming = opts?.streaming ?? true;
+		const messages: ChatCompletionInputMessage[] = opts?.messages ?? [
+			{ role: "user", content: "What is the capital of France?" },
+		];
+
+		const config = {
+			temperature: opts?.temperature,
+			max_tokens: opts?.max_tokens ?? 500,
+			top_p: opts?.top_p,
+		};
 		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": ${formatGenerationMessages({ messages, sep: ",\n    ", start: `[\n    `, end: `\n]` })},
+    ${formatGenerationConfig({ config, sep: ",\n    ", start: "", end: "" })},
+    "stream": ${!!streaming}
+}'`,
 		};
 	} else {
 		return snippetBasic(model, accessToken);
@@ -76,7 +120,7 @@ export const snippetFile = (model: ModelDataMinimal, accessToken: string): Infer
 export const curlSnippets: Partial<
 	Record<
 		PipelineType,
-		(model: ModelDataMinimal, accessToken: string, opts?: Record<string, string | boolean | number>) => InferenceSnippet
+		(model: ModelDataMinimal, accessToken: string, opts?: Record<string, unknown>) => InferenceSnippet
 	>
 > = {
 	// Same order as in js/src/lib/interfaces/Types.ts