Create configuration_protst.py

configuration_protst.py

@@ -0,0 +1,51 @@
+from transformers import PretrainedConfig
+from transformers.utils import logging
+from transformers.models.esm import EsmConfig
+from transformers.models.bert import BertConfig
+
+logger = logging.get_logger(__name__)
+
+
+class ProtSTConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`ProtSTModel`].
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        protein_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`EsmForProteinRepresentation`].
+        text_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`BertForPubMed`].
+    ```"""
+
+    model_type = "protst"
+
+    def __init__(
+        self,
+        protein_config=None,
+        text_config=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if protein_config is None:
+            protein_config = {}
+            logger.info("`protein_config` is `None`. Initializing the `ProtSTTextConfig` with default values.")
+
+        if text_config is None:
+            text_config = {}
+            logger.info("`text_config` is `None`. Initializing the `ProtSTVisionConfig` with default values.")
+
+        self.protein_config = EsmConfig(**protein_config)
+        self.text_config = BertConfig(**text_config)
+
+    @classmethod
+    def from_protein_text_configs(
+        cls, protein_config: EsmConfig, text_config: BertConfig, **kwargs
+    ):
+        r"""
+        Instantiate a [`ProtSTConfig`] (or a derived class) from ProtST text model configuration. Returns:
+            [`ProtSTConfig`]: An instance of a configuration object
+        """
+
+        return cls(protein_config=protein_config.to_dict(), text_config=text_config.to_dict(), **kwargs)