From 17e747a02d3553d90ef73d6c189734aef20c92e3 Mon Sep 17 00:00:00 2001
From: Mishig Davaadorj <dmishig@gmail.com>
Date: Mon, 4 Nov 2024 10:46:12 +0100
Subject: [PATCH] [transformers snippet] Support VLMs

---
 packages/tasks/src/library-to-tasks.ts        |  1 +
 .../tasks/src/model-libraries-snippets.ts     | 21 +++++++++++++++++--
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/packages/tasks/src/library-to-tasks.ts b/packages/tasks/src/library-to-tasks.ts
index c8411fbaab..1bcb910d46 100644
--- a/packages/tasks/src/library-to-tasks.ts
+++ b/packages/tasks/src/library-to-tasks.ts
@@ -55,6 +55,7 @@ export const LIBRARY_TASK_MAPPING: Partial<Record<ModelLibraryKey, PipelineType[
 		"image-segmentation",
 		"image-to-image",
 		"image-to-text",
+		"image-text-to-text",
 		"object-detection",
 		"question-answering",
 		"summarization",
diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts
index fc42805333..abeaf4da5a 100644
--- a/packages/tasks/src/model-libraries-snippets.ts
+++ b/packages/tasks/src/model-libraries-snippets.ts
@@ -849,11 +849,28 @@ export const transformers = (model: ModelData): string[] => {
 		const pipelineSnippet = ["# Use a pipeline as a high-level helper", "from transformers import pipeline", ""];
 
 		if (model.tags.includes("conversational") && model.config?.tokenizer_config?.chat_template) {
-			pipelineSnippet.push("messages = [", '    {"role": "user", "content": "Who are you?"},', "]");
+			if (model.pipeline_tag === "text-generation") {
+				pipelineSnippet.push("messages = [", '    {"role": "user", "content": "Who are you?"},', "]\n");
+			} else if (model.pipeline_tag === "image-text-to-text") {
+				pipelineSnippet.push(
+					`image_ny = "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg"`,
+					`image_chicago = "https://cdn.britannica.com/59/94459-050-DBA42467/Skyline-Chicago.jpg"`,
+					"\n"
+				);
+			}
 		}
 		pipelineSnippet.push(`pipe = pipeline("${model.pipeline_tag}", model="${model.id}"` + remote_code_snippet + ")");
 		if (model.tags.includes("conversational") && model.config?.tokenizer_config?.chat_template) {
-			pipelineSnippet.push("pipe(messages)");
+			if (model.pipeline_tag === "text-generation") {
+				pipelineSnippet.push("pipe(messages)");
+			} else if (model.pipeline_tag === "image-text-to-text") {
+				pipelineSnippet.push(
+					"pipe(",
+					"    images=[image_ny, image_chicago],",
+					`    text="<image> <image> Are these the same cities? If not what cities are these?",`,
+					")\n"
+				);
+			}
 		}
 
 		return [pipelineSnippet.join("\n"), autoSnippet];