Add API snippet zero-shot img classification (#363)

Author: osanseviero

packages/tasks/src/snippets/inputs.ts

@@ -79,6 +79,8 @@ const inputsTextToAudio = () => `"liquid drum and bass, atmospheric synths, airy
 
 const inputsAutomaticSpeechRecognition = () => `"sample1.flac"`;
 
+const inputsZeroShotImageClassification = () => `"cats.jpg"`;
+
 const modelInputSnippets: {
 	[key in PipelineType]?: (model: ModelData) => string;
 } = {
@@ -105,6 +107,7 @@ const modelInputSnippets: {
 	"token-classification": inputsTokenClassification,
 	translation: inputsTranslation,
 	"zero-shot-classification": inputsZeroShotClassification,
+	"zero-shot-image-classification": inputsZeroShotImageClassification,
 };
 
 // Use noWrap to put the whole snippet on a single line (removing new lines and tabulations)

packages/tasks/src/snippets/python.ts

@@ -12,6 +12,22 @@ output = query({
     "parameters": {"candidate_labels": ["refund", "legal", "faq"]},
 })`;
 
+export const snippetZeroShotImageClassification = (model: ModelData): string =>
+	`def query(data):
+	with open(data["image_path"], "rb") as f:
+		img = f.read()
+	payload={
+		"parameters": data["parameters"],
+		"inputs": base64.b64encode(img).decode("utf-8")
+	}
+	response = requests.post(API_URL, headers=headers, json=payload)
+	return response.json()
+
+output = query({
+    "image_path": ${getModelInputSnippet(model)},
+    "parameters": {"candidate_labels": ["cat", "dog", "llama"]},
+})`;
+
 export const snippetBasic = (model: ModelData): string =>
 	`def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
@@ -71,7 +87,7 @@ Audio(audio, rate=sampling_rate)`;
 	}
 };
 export const pythonSnippets: Partial<Record<PipelineType, (model: ModelData) => string>> = {
-	// Same order as in js/src/lib/interfaces/Types.ts
+	// Same order as in tasks/src/pipelines.ts
 	"text-classification": snippetBasic,
 	"token-classification": snippetBasic,
 	"table-question-answering": snippetBasic,
@@ -92,9 +108,10 @@ export const pythonSnippets: Partial<Record<PipelineType, (model: ModelData) =>
 	"audio-to-audio": snippetFile,
 	"audio-classification": snippetFile,
 	"image-classification": snippetFile,
-	"image-to-text": snippetFile,
 	"object-detection": snippetFile,
 	"image-segmentation": snippetFile,
+	"image-to-text": snippetFile,
+	"zero-shot-image-classification": snippetZeroShotImageClassification,
 };
 
 export function getPythonInferenceSnippet(model: ModelData, accessToken: string): string {