In "Use in transformers", use messages if conversational (#728)

Increase visibility of conversational capabilities of `text-generation` pipeline ![Screenshot from 2024-05-30 18-45-05](https://github.com/huggingface/huggingface.js/assets/7246357/e58af649-3943-49a8-815e-800c5d97aae0) --------- Co-authored-by: Julien Chaumond <[email protected]>
huggingface · May 31, 2024 · 8df8f6a · 8df8f6a
1 parent a352941
commit 8df8f6a
Showing 1 changed file with 11 additions and 7 deletions.
diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts
@@ -440,13 +440,17 @@ 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",
-			"",
-			`pipe = pipeline("${model.pipeline_tag}", model="${model.id}"` + remote_code_snippet + ")",
-		].join("\n");
-		return [pipelineSnippet, autoSnippet];
+		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?"},', "]");
+		}
+		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)");
+		}
+
+		return [pipelineSnippet.join("\n"), autoSnippet];
 	}
 	return [autoSnippet];
 };