diff --git a/gguf-py/scripts/gguf-template.py b/gguf-py/scripts/gguf-template.py index d3757b4f3..e5f80c69d 100644 --- a/gguf-py/scripts/gguf-template.py +++ b/gguf-py/scripts/gguf-template.py @@ -47,7 +47,9 @@ def get_chat_template(model_file: str, verbose: bool = False) -> str: return "" -def display_chat_template(chat_template: str, render_template: bool = False): +def display_chat_template( + chat_template: str, bos_token: str, eos_token: str, render_template: bool = False +): """ Display the chat template to standard output, optionally formatting it using Jinja2. @@ -68,9 +70,9 @@ def display_chat_template(chat_template: str, render_template: bool = False): {"role": "system", "content": "I am a helpful assistant."}, {"role": "user", "content": "Hello!"}, {"role": "assistant", "content": "Hello! How may I assist you today?"}, + {"role": "user", "content": "Can you tell me what pickled mayonnaise is?"}, + {"role": "assistant", "content": "Certainly! What would you like to know about it?"}, ] - bos_token = "" - eos_token = "" try: formatted_template = template.render( @@ -104,6 +106,18 @@ def main(): action="store_true", help="Render the chat template using Jinja2", ) + parser.add_argument( + "-b", + "--bos", + default="", + help="Set a bos special token. Default is ''.", + ) + parser.add_argument( + "-e", + "--eos", + default="", + help="Set a eos special token. Default is ''.", + ) parser.add_argument( "-v", "--verbose", @@ -114,7 +128,9 @@ def main(): model_file = args.model_file chat_template = get_chat_template(model_file, args.verbose) - display_chat_template(chat_template, render_template=args.render_template) + display_chat_template( + chat_template, args.bos, args.eos, render_template=args.render_template + ) if __name__ == "__main__":