KBaba7/llama.cpp

1#!/usr/bin/env python2'''3  Fetches the Jinja chat template of a HuggingFace model.4  If a model has multiple chat templates, you can specify the variant name.5 6  Syntax:7    ./scripts/get_chat_template.py model_id [variant]8 9  Examples:10    ./scripts/get_chat_template.py NousResearch/Meta-Llama-3-8B-Instruct11    ./scripts/get_chat_template.py NousResearch/Hermes-3-Llama-3.1-8B tool_use12    ./scripts/get_chat_template.py meta-llama/Llama-3.2-3B-Instruct13'''14 15import json16import re17import sys18 19 20def get_chat_template(model_id, variant=None):21    try:22        # Use huggingface_hub library if available.23        # Allows access to gated models if the user has access and ran `huggingface-cli login`.24        from huggingface_hub import hf_hub_download25        with open(hf_hub_download(repo_id=model_id, filename="tokenizer_config.json")) as f:26            config_str = f.read()27    except ImportError:28        import requests29        assert re.match(r"^[\w.-]+/[\w.-]+$", model_id), f"Invalid model ID: {model_id}"30        response = requests.get(f"https://huggingface.co/{model_id}/resolve/main/tokenizer_config.json")31        if response.status_code == 401:32            raise Exception('Access to this model is gated, please request access, authenticate with `huggingface-cli login` and make sure to run `pip install huggingface_hub`')33        response.raise_for_status()34        config_str = response.text35 36    try:37        config = json.loads(config_str)38    except json.JSONDecodeError:39        # Fix https://huggingface.co/NousResearch/Meta-Llama-3-8B-Instruct/blob/main/tokenizer_config.json40        # (Remove extra '}' near the end of the file)41        config = json.loads(re.sub(r'\}([\n\s]*\}[\n\s]*\],[\n\s]*"clean_up_tokenization_spaces")', r'\1', config_str))42 43    chat_template = config['chat_template']44    if isinstance(chat_template, str):45        return chat_template46    else:47        variants = {48            ct['name']: ct['template']49            for ct in chat_template50        }51 52        def format_variants():53            return ', '.join(f'"{v}"' for v in variants.keys())54 55        if variant is None:56            if 'default' not in variants:57                raise Exception(f'Please specify a chat template variant (one of {format_variants()})')58            variant = 'default'59            sys.stderr.write(f'Note: picked "default" chat template variant (out of {format_variants()})\n')60        elif variant not in variants:61            raise Exception(f"Variant {variant} not found in chat template (found {format_variants()})")62 63        return variants[variant]64 65 66def main(args):67    if len(args) < 1:68        raise ValueError("Please provide a model ID and an optional variant name")69    model_id = args[0]70    variant = None if len(args) < 2 else args[1]71 72    template = get_chat_template(model_id, variant)73    sys.stdout.write(template)74 75 76if __name__ == '__main__':77    main(sys.argv[1:])78