import os import json from transformers import PretrainedConfig, AutoConfig from utils.model_mapper import ModelMapper from typing import Optional, List, Dict, Any, Union from dataclasses import dataclass, field, asdict # model config class LlmConfig(PretrainedConfig): model_type = "llm_config" def __init__(self, **kwargs): self.hidden_size = kwargs.pop("hidden_size", 0) self.num_attention_heads = kwargs.pop("num_attention_heads", None) self.num_hidden_layers = kwargs.pop("num_hidden_layers", None) self.num_key_value_heads = kwargs.pop("num_key_value_heads", self.num_attention_heads) self.head_dim = kwargs.pop("head_dim", self.hidden_size // self.num_attention_heads if self.num_attention_heads else None) self.rope_theta = kwargs.pop("rope_theta", 10000.0) self.rope_ratio = kwargs.pop("rope_ratio", 1.0) self.sliding_window = kwargs.pop("sliding_window", 0) self.sliding_window = self.sliding_window if self.sliding_window is not None else 0 self.layer_types = kwargs.pop("layer_types", []) self.attention_type = kwargs.pop("attention_type", 'full') self.tie_word_embeddings = kwargs.pop("tie_word_embeddings", False) self.scale_emb = kwargs.pop("scale_emb", None) self.conv_L_cache = kwargs.pop("conv_L_cache", 0) self.rope_parameters = kwargs.pop("rope_parameters", None) self.qk_norm_after_rope = kwargs.pop("qk_norm_after_rope", False) self.model_map = kwargs.pop("model_map", {}) super().__init__(**kwargs) @staticmethod def _register_external_model(model_type: str): EXTERNAL_MODEL_REGISTRY = { 'funaudiochat': ('funaudiochat.register', 'register_funaudiochat'), } if model_type in EXTERNAL_MODEL_REGISTRY: module_path, func_name = EXTERNAL_MODEL_REGISTRY[model_type] try: import importlib module = importlib.import_module(module_path) getattr(module, func_name)() except ImportError: raise ImportError( f"{model_type} requires external package. " f"Please clone it from GitHub and set PYTHONPATH accordingly." ) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): config_path = os.path.join(pretrained_model_name_or_path, 'config.json') if os.path.exists(config_path): with open(config_path, 'r') as f: raw_config = json.load(f) model_type = raw_config.get('model_type') cls._register_external_model(model_type) # Handle models without top-level model_type (e.g., lfm2_audio) if model_type is None: archs = raw_config.get('architectures', []) if 'Lfm2AudioForConditionalGeneration' in archs and 'lfm' in raw_config: from transformers import Lfm2Config config = Lfm2Config(**raw_config['lfm']) config.model_type = 'lfm2_audio' else: config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True, **kwargs) else: config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True, **kwargs) model_type, model_map = ModelMapper().get_map(config) llm_config_kwargs = { 'origin_config': config, 'model_type': model_type, 'model_map': model_map } llm_config = cls(**llm_config_kwargs) # rename attribute for different models ModelMapper.do_map(llm_config, config, model_map['config']) # Post-processing and setting defaults if llm_config.num_key_value_heads is None: llm_config.num_key_value_heads = llm_config.num_attention_heads # Compatibility: transformers>=5.x moved rope_theta into rope_parameters dict if llm_config.rope_theta is None or llm_config.rope_theta == 10000.0: # Try rope_parameters (transformers 5.x style) rp = getattr(config, 'rope_parameters', None) or getattr(config, 'rope_scaling', None) if isinstance(rp, dict) and 'rope_theta' in rp: llm_config.rope_theta = rp['rope_theta'] # Fallback to raw config JSON elif 'rope_theta' in raw_config: llm_config.rope_theta = raw_config['rope_theta'] if llm_config.rope_theta is None: llm_config.rope_theta = 10000.0 if llm_config.rope_ratio is None: llm_config.rope_ratio = 1.0 if llm_config.head_dim is None and llm_config.hidden_size and llm_config.num_attention_heads: if isinstance(llm_config.num_attention_heads, list): llm_config.head_dim = [llm_config.hidden_size // atten_head for atten_head in llm_config.num_attention_heads] else: llm_config.head_dim = llm_config.hidden_size // llm_config.num_attention_heads # Determine attention type. # Only `sliding_attention` layers need a sliding-window mask and # therefore the "mix" path. `linear_attention` layers (e.g. Qwen3.5) # do not use the attention mask at all, so they are treated as full. sliding_attn_layers = [] if hasattr(llm_config, 'layer_types') and llm_config.layer_types: for i in range(len(llm_config.layer_types)): if llm_config.layer_types[i] == 'sliding_attention': sliding_attn_layers.append(i) if llm_config.num_hidden_layers and len(sliding_attn_layers) >= llm_config.num_hidden_layers: llm_config.attention_type = 'sliding' elif len(sliding_attn_layers) > 0: llm_config.attention_type = 'mix' llm_config.sliding_attn_layers = sliding_attn_layers else: llm_config.attention_type = 'full' return llm_config # export config @dataclass class VisionExportConfig: """Configuration for vision-related capabilities.""" image_mean: Optional[List[float]] = field(default_factory=list) image_norm: Optional[List[float]] = field(default_factory=list) image_size: Optional[Union[int, List[int]]] = None image_size_unit: Optional[int] = None vision_start: Optional[int] = None vision_end: Optional[int] = None image_pad: Optional[int] = None num_grid_per_side: Optional[int] = None has_deepstack: bool = False image_max_size: Optional[int] = None global_image: Optional[int] = None vision_id_start_id: Optional[int] = None vision_id_end_id: Optional[int] = None vision_slice_start_id: Optional[int] = None vision_slice_end_id: Optional[int] = None @dataclass class LLMExportConfig: """Top-level container for all export configurations.""" is_audio: bool = False is_visual: bool = False has_talker: bool = False attention_mask: str = 'float' attention_type: str = 'full' sliding_window: int = 0 tie_embeddings: Optional[Union[List[int], Dict[str, Any]]] = field(default_factory=list) jinja: Dict[str, Any] = field(default_factory=dict) vision: Optional[VisionExportConfig] = None def to_dict(self) -> Dict[str, Any]: """Converts the configuration to a dictionary for JSON serialization.""" nested_dict = asdict(self) vision_data = nested_dict.pop('vision', None) if vision_data: nested_dict.update(vision_data) final_dict = {key: value for key, value in nested_dict.items() if value is not None} return final_dict