142 lines
5.1 KiB
Python
142 lines
5.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapted from SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/hf_transformers_utils.py
|
|
|
|
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Utilities for Huggingface Transformers."""
|
|
|
|
import contextlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from huggingface_hub import snapshot_download
|
|
from transformers import AutoConfig, PretrainedConfig
|
|
from transformers.models.auto.modeling_auto import (MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
|
|
|
|
_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = {
|
|
# ChatGLMConfig.model_type: ChatGLMConfig,
|
|
# DbrxConfig.model_type: DbrxConfig,
|
|
# ExaoneConfig.model_type: ExaoneConfig,
|
|
# Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig,
|
|
}
|
|
|
|
for name, cls in _CONFIG_REGISTRY.items():
|
|
with contextlib.suppress(ValueError):
|
|
AutoConfig.register(name, cls)
|
|
|
|
|
|
def download_from_hf(model_path: str):
|
|
if os.path.exists(model_path):
|
|
return model_path
|
|
|
|
return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"])
|
|
|
|
|
|
def get_hf_config(
|
|
model: str,
|
|
trust_remote_code: bool,
|
|
revision: str | None = None,
|
|
model_override_args: dict | None = None,
|
|
**kwargs,
|
|
):
|
|
is_gguf = check_gguf_file(model)
|
|
if is_gguf:
|
|
raise NotImplementedError("GGUF models are not supported.")
|
|
|
|
config = AutoConfig.from_pretrained(model, trust_remote_code=trust_remote_code, revision=revision, **kwargs)
|
|
if config.model_type in _CONFIG_REGISTRY:
|
|
config_class = _CONFIG_REGISTRY[config.model_type]
|
|
config = config_class.from_pretrained(model, revision=revision)
|
|
# NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`.
|
|
config._name_or_path = model
|
|
if model_override_args:
|
|
config.update(model_override_args)
|
|
|
|
# Special architecture mapping check for GGUF models
|
|
if is_gguf:
|
|
if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
|
|
raise RuntimeError(f"Can't get gguf config for {config.model_type}.")
|
|
model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type]
|
|
config.update({"architectures": [model_type]})
|
|
|
|
return config
|
|
|
|
|
|
def get_diffusers_config(model: str, ) -> dict[str, Any]:
|
|
"""Gets a configuration for the given diffusers model.
|
|
|
|
Args:
|
|
model: The model name or path.
|
|
fastvideo_args: Optional inference arguments to override in the config.
|
|
|
|
Returns:
|
|
The loaded configuration.
|
|
"""
|
|
config_name = "config.json"
|
|
if "scheduler" in model:
|
|
config_name = "scheduler_config.json"
|
|
# Check if the model path exists
|
|
if os.path.exists(model):
|
|
config_file = os.path.join(model, config_name)
|
|
if os.path.exists(config_file):
|
|
try:
|
|
# Load the config directly from the file
|
|
with open(config_file) as f:
|
|
config_dict: dict[str, Any] = json.load(f)
|
|
if "_diffusers_version" in config_dict:
|
|
config_dict.pop("_diffusers_version")
|
|
# TODO(will): apply any overrides from inference args
|
|
return config_dict
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load diffusers config from {config_file}: {e}") from e
|
|
raise RuntimeError(f"Config file not found at {config_file}")
|
|
else:
|
|
raise RuntimeError(f"Diffusers config file not found at {model}")
|
|
|
|
|
|
# Models don't use the same configuration key for determining the maximum
|
|
# context length. Store them here so we can sanely check them.
|
|
# NOTE: The ordering here is important. Some models have two of these and we
|
|
# have a preference for which value gets used.
|
|
CONTEXT_LENGTH_KEYS = [
|
|
"max_sequence_length",
|
|
"seq_length",
|
|
"max_seq_len",
|
|
"model_max_length",
|
|
"max_position_embeddings",
|
|
]
|
|
|
|
|
|
def attach_additional_stop_token_ids(tokenizer):
|
|
# Special handling for stop token <|eom_id|> generated by llama 3 tool use.
|
|
if "<|eom_id|>" in tokenizer.get_added_vocab():
|
|
tokenizer.additional_stop_token_ids = set([tokenizer.get_added_vocab()["<|eom_id|>"]])
|
|
else:
|
|
tokenizer.additional_stop_token_ids = None
|
|
|
|
|
|
def check_gguf_file(model: str | os.PathLike) -> bool:
|
|
"""Check if the file is a GGUF model."""
|
|
model = Path(model)
|
|
if not model.is_file():
|
|
return False
|
|
elif model.suffix == ".gguf":
|
|
return True
|
|
|
|
with open(model, "rb") as f:
|
|
header = f.read(4)
|
|
return header == b"GGUF"
|