Megrez-3B-Omni / configuration_megrezo.py
zhoudong
modify_license
483b1a4
# Copyright 2024 Infinigence AI Team. All rights reserved.
#
# 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.
"""MegrezO model configuration"""
from typing import Optional
from transformers.configuration_utils import PretrainedConfig
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.utils import logging
from .modeling_navit_siglip import SiglipVisionConfig
logger = logging.get_logger(__name__)
class AudioConfig(PretrainedConfig):
model_type = "megrezo"
def __init__(
self,
n_mels: int = 128,
n_ctx: int = 1500,
n_state: int = 1280,
n_head: int = 20,
n_layer: int = 32,
output_dim: int = 2560,
avg_pool: bool = True,
add_audio_bos_eos_token: bool = True,
**kwargs,
):
super().__init__(**kwargs)
self.n_mels = n_mels
self.n_ctx = n_ctx
self.n_state = n_state
self.n_head = n_head
self.n_layer = n_layer
self.output_dim = output_dim
self.avg_pool = avg_pool
self.add_audio_bos_eos_token = add_audio_bos_eos_token
class MegrezOConfig(LlamaConfig):
model_type = "megrezo"
keys_to_ignore_at_inference = ["past_key_values"]
is_composition = True
_default_audio_config = {
"n_mels": 128,
"n_ctx": 1500,
"n_state": 1280,
"n_head": 20,
"n_layer": 32,
"output_dim": 2560,
"avg_pool": True,
"add_audio_bos_eos_token": True,
}
_default_vision_config = {
"intermediate_size": 4304,
"num_hidden_layers": 27,
"num_attention_heads": 16,
"image_size": 980,
"hidden_size": 1152,
"patch_size": 16,
"model_type": "siglip_vision_model",
}
def __init__(
self,
audio_config: Optional[AudioConfig] = None,
vision_config: Optional[SiglipVisionConfig] = None,
**kwargs,
):
super().__init__(**kwargs)
if audio_config is None:
self.audio_config = AudioConfig(**self._default_audio_config)
elif isinstance(audio_config, dict):
self.audio_config = AudioConfig(**audio_config)
elif isinstance(audio_config, AudioConfig):
self.audio_config = audio_config
if vision_config is None:
self.vision_config = SiglipVisionConfig(**self._default_vision_config)
elif isinstance(vision_config, dict):
self.vision_config = SiglipVisionConfig(**vision_config)
elif isinstance(vision_config, SiglipVisionConfig):
self.vision_config = vision_config