Add script to check inference snippet changes + update Python textToImage snippet

Author: Wauplin

.github/workflows/inference-check-snippets.yml

@@ -0,0 +1,49 @@
+name: Inference check snippets
+on:
+  pull_request:
+    paths:
+      - "packages/tasks/src/snippets/**"
+      - ".github/workflows/inference-check-snippets.yml"
+
+jobs:
+  check-snippets:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - run: corepack enable
+
+      - uses: actions/setup-node@v3
+        with:
+          node-version: "20"
+          cache: "pnpm"
+          cache-dependency-path: "**/pnpm-lock.yaml"
+      - run: |
+          cd packages/tasks
+          pnpm install --frozen-lockfile --filter .
+          pnpm install --frozen-lockfile --filter ...[${{ steps.since.outputs.SINCE }}]...
+          pnpm --filter ...[${{ steps.since.outputs.SINCE }}]... build
+
+      # TODO: Find a way to run on all pipeline tags
+      # TODO: print snippet only if it has changed since the last commit on main (?)
+      # TODO: (even better: automated message on the PR with diff)
+      - name: Print text-to-image snippets
+        run: |
+          cd packages/tasks
+          pnpm run check-snippets --pipeline-tag="text-to-image"
+
+      - name: Print simple text-generation snippets
+        run: |
+          cd packages/tasks
+          pnpm run check-snippets --pipeline-tag="text-generation"
+
+      - name: Print conversational text-generation snippets
+        run: |
+          cd packages/tasks
+          pnpm run check-snippets --pipeline-tag="text-generation" --tags="conversational"
+
+      - name: Print conversational image-text-to-text snippets
+        run: |
+          cd packages/tasks
+          pnpm run check-snippets --pipeline-tag="image-text-to-text" --tags="conversational"

packages/tasks/scripts/check-snippets.ts

@@ -17,7 +17,7 @@ import minimist from "minimist";
 const args = minimist(process.argv.slice(2));
 
 const accessToken = "hf_**********";
-const pipelineTag = args["pipeline-type"] || "text-generation";
+const pipelineTag = args["pipeline-tag"] || "text-generation";
 const tags = (args["tags"] || "").split(",");
 
 const modelMinimal: ModelDataMinimal = {

packages/tasks/src/snippets/python.ts

@@ -4,6 +4,11 @@ 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}"}")`;
+
 export const snippetConversational = (
 	model: ModelDataMinimal,
 	accessToken: string,
@@ -184,18 +189,31 @@ export const snippetFile = (model: ModelDataMinimal): InferenceSnippet => ({
 output = query(${getModelInputSnippet(model)})`,
 });
 
-export const snippetTextToImage = (model: ModelDataMinimal): InferenceSnippet => ({
-	content: `def query(payload):
+export const snippetTextToImage = (model: ModelDataMinimal, accessToken: string): InferenceSnippet[] => {
+	return [
+		{
+			client: "requests",
+			content: `def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
 	return response.content
+
 image_bytes = query({
 	"inputs": ${getModelInputSnippet(model)},
 })
 # You can access the image with PIL.Image for example
 import io
 from PIL import Image
 image = Image.open(io.BytesIO(image_bytes))`,
-});
+		},
+		{
+			client: "huggingface_hub",
+			content: `${snippetImportInferenceClient(model, accessToken)}
+
+# output is a PIL.Image object
+image = client.text_to_image(${getModelInputSnippet(model)})`,
+		},
+	];
+};
 
 export const snippetTabular = (model: ModelDataMinimal): InferenceSnippet => ({
 	content: `def query(payload):
@@ -300,6 +318,9 @@ export function getPythonInferenceSnippet(
 	if (model.tags.includes("conversational")) {
 		// Conversational model detected, so we display a code snippet that features the Messages API
 		return snippetConversational(model, accessToken, opts);
+	} else if (model.pipeline_tag == "text-to-image") {
+		// TODO: factorize this logic
+		return snippetTextToImage(model, accessToken);
 	} else {
 		let snippets =
 			model.pipeline_tag && model.pipeline_tag in pythonSnippets