| """Configuration for pruned FlexOlmo models with variable-width expert 1.""" | |
| from transformers import FlexOlmoConfig | |
| class PrunedFlexOlmoConfig(FlexOlmoConfig): | |
| """Config for FlexOlmo with a pruned expert 1. | |
| Extends FlexOlmoConfig with expert_1_intermediate_size to specify | |
| the width of the pruned expert. | |
| """ | |
| model_type = "pruned_flex_olmo" | |
| def __init__(self, expert_1_intermediate_size: int = None, **kwargs): | |
| super().__init__(**kwargs) | |
| # expert_1_intermediate_size: width of pruned expert 1 | |
| # If None, falls back to intermediate_size (no pruning) | |
| self.expert_1_intermediate_size = expert_1_intermediate_size or self.intermediate_size | |