| """ | |
| Configuration class for MoLA-LM | |
| """ | |
| from transformers import PretrainedConfig | |
| from typing import Dict, List | |
| EXPERTS_LIST = [ | |
| "0", | |
| "1", | |
| "2", | |
| "3", | |
| "4", | |
| "5", | |
| "6", | |
| "7", | |
| "8", | |
| ] | |
| class MoLAConfig(PretrainedConfig): | |
| """Configuration class for MoLA-LM model.""" | |
| model_type = "mola_lm" | |
| def __init__( | |
| self, | |
| base_model_name_or_path: str = "Qwen/Qwen2.5-3B-Instruct", | |
| task_labels: List[str] = None, | |
| router_config: Dict = None, | |
| lora_configs: Dict[str, Dict] = None, | |
| **kwargs | |
| ): | |
| super().__init__(**kwargs) | |
| self.base_model_name_or_path = base_model_name_or_path | |
| self.task_labels = task_labels or EXPERTS_LIST | |
| self.router_config = router_config or {} | |
| self.lora_configs = lora_configs or {} | |
| self.num_loras = len(self.task_labels) | |