diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts index ed146b409b..19e7ce666a 100644 --- a/packages/tasks/src/model-libraries-snippets.ts +++ b/packages/tasks/src/model-libraries-snippets.ts @@ -1142,13 +1142,31 @@ export const transformers = (model: ModelData): string[] => { } if (model.pipeline_tag && LIBRARY_TASK_MAPPING.transformers?.includes(model.pipeline_tag)) { - const pipelineSnippet = ["# Use a pipeline as a high-level helper", "from transformers import pipeline", ""]; + const pipelineSnippet = [ + "# Use a pipeline as a high-level helper", + "from transformers import pipeline", + "", + `pipe = pipeline("${model.pipeline_tag}", model="${model.id}"` + remote_code_snippet + ")", + ]; - if (model.tags.includes("conversational") && model.config?.tokenizer_config?.chat_template) { - pipelineSnippet.push("messages = [", ' {"role": "user", "content": "Who are you?"},', "]"); - } - pipelineSnippet.push(`pipe = pipeline("${model.pipeline_tag}", model="${model.id}"` + remote_code_snippet + ")"); - if (model.tags.includes("conversational") && model.config?.tokenizer_config?.chat_template) { + if (model.tags.includes("conversational")) { + if (model.tags.includes("image-text-to-text")) { + pipelineSnippet.push( + "messages = [", + [ + " {", + ' "role": "user",', + ' "content": [', + ' {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"},', + ' {"type": "text", "text": "What animal is on the candy?"}', + " ]", + " },", + ].join("\n"), + "]" + ); + } else { + pipelineSnippet.push("messages = [", ' {"role": "user", "content": "Who are you?"},', "]"); + } pipelineSnippet.push("pipe(messages)"); }