| # coding=utf-8 | |
| # Copyright 2024 MultiomicsLM 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. | |
| """MultiomicsLM ESM model configuration.""" | |
| from transformers.models.esm.configuration_esm import EsmConfig | |
| from transformers.configuration_utils import PretrainedConfig | |
| from typing import Optional | |
| class MMEsmConfig(EsmConfig): | |
| """ | |
| Configuration class for MultiomicsLM ESM model. | |
| This configuration extends the base ESM configuration with multiomics-specific parameters. | |
| """ | |
| model_type = "esm" | |
| num_labels = 2 | |
| def __init__( | |
| self, | |
| structure: Optional[dict] = None, | |
| **kwargs | |
| ): | |
| super().__init__(**kwargs) | |
| default_structure = { | |
| "width": 512, | |
| "projector": "mlp3x_gelu", | |
| "output_dim": self.hidden_size if hasattr(self, "hidden_size") else 640, | |
| "projector_norm": True, | |
| "embedding_keys": ["esm_if_emb"], | |
| } | |
| if structure is None: | |
| self.structure = default_structure | |
| else: | |
| cfg = dict(default_structure) | |
| cfg.update(structure) | |
| self.structure = cfg | |