| 17 | |
| 18 | |
| 19 | def get_chat_template(model_id, variant=None): |
| 20 | try: |
| 21 | # Use huggingface_hub library if available. |
| 22 | # Allows access to gated models if the user has access and ran `huggingface-cli login`. |
| 23 | from huggingface_hub import hf_hub_download |
| 24 | with open(hf_hub_download(repo_id=model_id, filename="tokenizer_config.json"), encoding="utf-8") as f: |
| 25 | config_str = f.read() |
| 26 | except ImportError: |
| 27 | import requests |
| 28 | assert re.match(r"^[\w.-]+/[\w.-]+$", model_id), f"Invalid model ID: {model_id}" |
| 29 | response = requests.get(f"https://huggingface.co/{model_id}/resolve/main/tokenizer_config.json") |
| 30 | if response.status_code == 401: |
| 31 | 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`') |
| 32 | response.raise_for_status() |
| 33 | config_str = response.text |
| 34 | |
| 35 | try: |
| 36 | config = json.loads(config_str) |
| 37 | except json.JSONDecodeError: |
| 38 | # Fix https://huggingface.co/NousResearch/Meta-Llama-3-8B-Instruct/blob/main/tokenizer_config.json |
| 39 | # (Remove extra '}' near the end of the file) |
| 40 | config = json.loads(re.sub(r'\}([\n\s]*\}[\n\s]*\],[\n\s]*"clean_up_tokenization_spaces")', r'\1', config_str)) |
| 41 | |
| 42 | chat_template = config['chat_template'] |
| 43 | if isinstance(chat_template, str): |
| 44 | return chat_template |
| 45 | else: |
| 46 | variants = { |
| 47 | ct['name']: ct['template'] |
| 48 | for ct in chat_template |
| 49 | } |
| 50 | |
| 51 | def format_variants(): |
| 52 | return ', '.join(f'"{v}"' for v in variants.keys()) |
| 53 | |
| 54 | if variant is None: |
| 55 | if 'default' not in variants: |
| 56 | raise Exception(f'Please specify a chat template variant (one of {format_variants()})') |
| 57 | variant = 'default' |
| 58 | sys.stderr.write(f'Note: picked "default" chat template variant (out of {format_variants()})\n') |
| 59 | elif variant not in variants: |
| 60 | raise Exception(f"Variant {variant} not found in chat template (found {format_variants()})") |
| 61 | |
| 62 | return variants[variant] |
| 63 | |
| 64 | |
| 65 | def main(args): |