Skip to content

Commit 8df8f6a

Browse files
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]>
1 parent a352941 commit 8df8f6a

File tree

1 file changed

+11
-7
lines changed

1 file changed

+11
-7
lines changed

packages/tasks/src/model-libraries-snippets.ts

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -440,13 +440,17 @@ export const transformers = (model: ModelData): string[] => {
440440
}
441441

442442
if (model.pipeline_tag && LIBRARY_TASK_MAPPING.transformers?.includes(model.pipeline_tag)) {
443-
const pipelineSnippet = [
444-
"# Use a pipeline as a high-level helper",
445-
"from transformers import pipeline",
446-
"",
447-
`pipe = pipeline("${model.pipeline_tag}", model="${model.id}"` + remote_code_snippet + ")",
448-
].join("\n");
449-
return [pipelineSnippet, autoSnippet];
443+
const pipelineSnippet = ["# Use a pipeline as a high-level helper", "from transformers import pipeline", ""];
444+
445+
if (model.tags.includes("conversational") && model.config?.tokenizer_config?.chat_template) {
446+
pipelineSnippet.push("messages = [", ' {"role": "user", "content": "Who are you?"},', "]");
447+
}
448+
pipelineSnippet.push(`pipe = pipeline("${model.pipeline_tag}", model="${model.id}"` + remote_code_snippet + ")");
449+
if (model.tags.includes("conversational") && model.config?.tokenizer_config?.chat_template) {
450+
pipelineSnippet.push("pipe(messages)");
451+
}
452+
453+
return [pipelineSnippet.join("\n"), autoSnippet];
450454
}
451455
return [autoSnippet];
452456
};

0 commit comments

Comments
 (0)