Upload configuration_esm.py

configuration_esm.py

@@ -0,0 +1,51 @@
+# 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
+