add support of RnaMsm

Signed-off-by: Zhiyuan Chen <[email protected]>

multimolecule/__init__.py

@@ -6,17 +6,29 @@
     RnaBertForSequenceClassification,
     RnaBertForTokenClassification,
     RnaBertModel,
+    RnaMsmConfig,
+    RnaMsmForMaskedLM,
+    RnaMsmForPretraining,
+    RnaMsmForSequenceClassification,
+    RnaMsmForTokenClassification,
+    RnaMsmModel,
 )
 from .tokenizers import RnaTokenizer
 
 __all__ = [
     "models",
     "tokenizers",
+    "RnaTokenizer",
     "RnaBertConfig",
     "RnaBertModel",
     "RnaBertForMaskedLM",
     "RnaBertForPretraining",
     "RnaBertForSequenceClassification",
     "RnaBertForTokenClassification",
-    "RnaTokenizer",
+    "RnaMsmConfig",
+    "RnaMsmModel",
+    "RnaMsmForMaskedLM",
+    "RnaMsmForPretraining",
+    "RnaMsmForSequenceClassification",
+    "RnaMsmForTokenClassification",
 ]

multimolecule/models/__init__.py

@@ -1,19 +1,33 @@
+from ..tokenizers.rna import RnaTokenizer
 from .rnabert import (
     RnaBertConfig,
     RnaBertForMaskedLM,
     RnaBertForPretraining,
     RnaBertForSequenceClassification,
     RnaBertForTokenClassification,
     RnaBertModel,
-    RnaTokenizer,
+)
+from .rnamsm import (
+    RnaMsmConfig,
+    RnaMsmForMaskedLM,
+    RnaMsmForPretraining,
+    RnaMsmForSequenceClassification,
+    RnaMsmForTokenClassification,
+    RnaMsmModel,
 )
 
 __all__ = [
+    "RnaTokenizer",
     "RnaBertConfig",
     "RnaBertModel",
     "RnaBertForMaskedLM",
     "RnaBertForPretraining",
     "RnaBertForSequenceClassification",
     "RnaBertForTokenClassification",
-    "RnaTokenizer",
+    "RnaMsmConfig",
+    "RnaMsmModel",
+    "RnaMsmForMaskedLM",
+    "RnaMsmForPretraining",
+    "RnaMsmForSequenceClassification",
+    "RnaMsmForTokenClassification",
 ]

multimolecule/models/rnamsm/__init__.py

@@ -0,0 +1,40 @@
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForMaskedLM,
+    AutoModelForPreTraining,
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+    AutoModelWithLMHead,
+    AutoTokenizer,
+)
+
+from multimolecule.tokenizers.rna import RnaTokenizer
+
+from .configuration_rnamsm import RnaMsmConfig
+from .modeling_rnamsm import (
+    RnaMsmForMaskedLM,
+    RnaMsmForPretraining,
+    RnaMsmForSequenceClassification,
+    RnaMsmForTokenClassification,
+    RnaMsmModel,
+)
+
+__all__ = [
+    "RnaMsmConfig",
+    "RnaMsmModel",
+    "RnaTokenizer",
+    "RnaMsmForPretraining",
+    "RnaMsmForMaskedLM",
+    "RnaMsmForSequenceClassification",
+    "RnaMsmForTokenClassification",
+]
+
+AutoConfig.register("rnamsm", RnaMsmConfig)
+AutoModel.register(RnaMsmConfig, RnaMsmModel)
+AutoModelForPreTraining.register(RnaMsmConfig, RnaMsmForPretraining)
+AutoModelForMaskedLM.register(RnaMsmConfig, RnaMsmForMaskedLM)
+AutoModelForSequenceClassification.register(RnaMsmConfig, RnaMsmForSequenceClassification)
+AutoModelForTokenClassification.register(RnaMsmConfig, RnaMsmForTokenClassification)
+AutoModelWithLMHead.register(RnaMsmConfig, RnaMsmForTokenClassification)
+AutoTokenizer.register(RnaMsmConfig, RnaTokenizer)

multimolecule/models/rnamsm/configuration_rnamsm.py

@@ -0,0 +1,123 @@
+from transformers.utils import logging
+
+from ..configuration_utils import HeadConfig, MaskedLMHeadConfig, PretrainedConfig
+
+logger = logging.get_logger(__name__)
+
+
+class RnaMsmConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`RnaMsmModel`]. It is used to instantiate a
+    RnaMsm model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the RnaMsm
+    [yikunpku/RNA-MSM](https://github.com/yikunpku/RNA-MSM) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*):
+            Vocabulary size of the RnaMsm model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`RnaMsmModel`].
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
+        hidden_dropout (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_dropout (`float`, *optional*, defaults to 0.1):
+            The dropout ratio for the attention probabilities.
+        max_position_embeddings (`int`, *optional*, defaults to 1026):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        pad_token_id (`int`, *optional*, defaults to 0):
+            The index of the padding token in the vocabulary. This must be included in the config because certain parts
+            of the RnaBert code use this instead of the attention mask.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            The index of the bos token in the vocabulary. This must be included in the config because of the
+            contact and other prediction heads removes the bos and padding token when predicting outputs.
+        mask_token_id (`int`, *optional*, defaults to 4):
+            The index of the mask token in the vocabulary. This must be included in the config because of the
+            "mask-dropout" scaling trick, which will scale the inputs depending on the number of masked tokens.
+
+    Examples:
+
+    ```python
+    >>> from multimolecule import RnaMsmModel, RnaMsmConfig
+
+    >>> # Initializing a RnaMsm style configuration >>> configuration = RnaMsmConfig()
+
+    >>> # Initializing a model from the configuration >>> model = RnaMsmModel(configuration)
+
+    >>> # Accessing the model configuration >>> configuration = model.config
+    ```
+    """
+
+    model_type = "rnamsm"
+
+    def __init__(
+        self,
+        vocab_size=25,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout=0.1,
+        attention_dropout=0.1,
+        max_position_embeddings=1024,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        pad_token_id=0,
+        bos_token_id=1,
+        mask_token_id=4,
+        position_embedding_type="absolute",
+        use_cache=True,
+        max_tokens_per_msa=2**14,
+        attention_type="standard",
+        embed_positions_msa=True,
+        attention_bias=True,
+        head=None,
+        lm_head=None,
+        **kwargs,
+    ):
+        if head is None:
+            head = {}
+        head.setdefault("hidden_size", hidden_size)
+        if "problem_type" in kwargs:
+            head.setdefault("problem_type", kwargs["problem_type"])
+        if "num_labels" in kwargs:
+            head.setdefault("num_labels", kwargs["num_labels"])
+        if lm_head is None:
+            lm_head = {}
+        lm_head.setdefault("hidden_size", hidden_size)
+        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, mask_token_id=mask_token_id, **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout = hidden_dropout
+        self.attention_dropout = attention_dropout
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.position_embedding_type = position_embedding_type
+        self.use_cache = use_cache
+        self.max_tokens_per_msa = max_tokens_per_msa
+        self.attention_type = attention_type
+        self.embed_positions_msa = embed_positions_msa
+        self.attention_bias = attention_bias
+        self.head = HeadConfig(**head)
+        self.lm_head = MaskedLMHeadConfig(**lm_head)

multimolecule/models/rnamsm/convert_checkpoint.py

@@ -0,0 +1,117 @@
+import os
+from typing import Optional
+
+import chanfig
+import torch
+from torch import nn
+
+from multimolecule.models import RnaMsmConfig as Config
+from multimolecule.models import RnaMsmForMaskedLM as Model
+from multimolecule.tokenizers.rna.utils import get_special_tokens_map, get_tokenizer_config, get_vocab_list
+
+try:
+    from huggingface_hub import HfApi
+except ImportError:
+    HfApi = None
+
+
+torch.manual_seed(1013)
+
+CONFIG = {
+    "architectures": ["RnaMsmModel"],
+    "attention_dropout": 0.1,
+    "hidden_act": "gelu",
+    "hidden_dropout": 0.1,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "max_position_embeddings": 1024,
+    "num_attention_heads": 12,
+    "num_hidden_layers": 10,
+    "vocab_size": 25,
+    "num_labels": 1,
+    "pad_token_id": 0,
+    "embed_positions_msa": True,
+}
+
+original_vocab_list = ["<cls>", "<pad>", "<eos>", "<unk>", "A", "G", "C", "U", "X", "N", "-", "<mask>"]
+vocab_list = get_vocab_list()
+
+
+def _convert_checkpoint(config, original_state_dict):
+    state_dict = {}
+    for key, value in original_state_dict.items():
+        key = key.replace("layers", "rnamsm.encoder.layer")
+        key = key.replace("msa_position_embedding", "rnamsm.embeddings.msa_embeddings")
+        key = key.replace("embed_tokens", "rnamsm.embeddings.word_embeddings")
+        key = key.replace("embed_positions", "rnamsm.embeddings.position_embeddings")
+        key = key.replace("emb_layer_norm_before", "rnamsm.embeddings.layer_norm")
+        key = key.replace("emb_layer_norm_after", "rnamsm.encoder.layer_norm")
+        key = key.replace("regression", "decoder")
+        key = key.replace("lm_head.dense", "lm_head.transform.dense")
+        key = key.replace("lm_head.layer_norm", "lm_head.transform.layer_norm")
+        state_dict[key] = value
+
+    state_vocab_size = state_dict["rnamsm.embeddings.word_embeddings.weight"].size(0)
+    original_vocab_size = len(original_vocab_list)
+    if state_vocab_size != original_vocab_size:
+        raise ValueError(
+            f"Vocabulary size do not match. Expected to have {original_vocab_size}, but got {state_vocab_size}."
+        )
+    word_embed = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+    word_embed_weight = word_embed.weight.data
+    predictions_bias = torch.zeros(config.vocab_size)
+    # nn.init.normal_(pos_embed.weight, std=0.02)
+    for original_index, original_token in enumerate(original_vocab_list):
+        new_index = vocab_list.index(original_token)
+        word_embed_weight[new_index] = state_dict["rnamsm.embeddings.word_embeddings.weight"][original_index]
+        predictions_bias[new_index] = state_dict["lm_head.bias"][original_index]
+    state_dict["rnamsm.embeddings.word_embeddings.weight"] = state_dict["lm_head.decoder.weight"] = word_embed_weight
+    state_dict["lm_head.decoder.bias"] = state_dict["lm_head.bias"] = predictions_bias
+    del state_dict["lm_head.weight"]
+    return state_dict
+
+
+def convert_checkpoint(convert_config):
+    config = Config.from_dict(chanfig.FlatDict(CONFIG))
+    config.vocab_size = len(vocab_list)
+
+    model = Model(config)
+
+    ckpt = torch.load(convert_config.checkpoint_path, map_location=torch.device("cpu"))
+    state_dict = _convert_checkpoint(config, ckpt)
+
+    model.load_state_dict(state_dict)
+    model.save_pretrained(convert_config.output_path, safe_serialization=True)
+    model.save_pretrained(convert_config.output_path, safe_serialization=False)
+    chanfig.NestedDict(get_special_tokens_map()).json(
+        os.path.join(convert_config.output_path, "special_tokens_map.json")
+    )
+    chanfig.NestedDict(get_tokenizer_config()).json(os.path.join(convert_config.output_path, "tokenizer_config.json"))
+
+    if convert_config.push_to_hub:
+        if HfApi is None:
+            raise ImportError("Please install huggingface_hub to push to the hub.")
+        api = HfApi()
+        api.create_repo(
+            convert_config.repo_id,
+            token=convert_config.token,
+            exist_ok=True,
+        )
+        api.upload_folder(
+            repo_id=convert_config.repo_id, folder_path=convert_config.output_path, token=convert_config.token
+        )
+
+
+@chanfig.configclass
+class ConvertConfig:
+    checkpoint_path: str
+    output_path: str = Config.model_type
+    push_to_hub: bool = False
+    repo_id: str = f"multimolecule/{output_path}"
+    token: Optional[str] = None
+
+
+if __name__ == "__main__":
+    config = ConvertConfig()
+    config.parse()  # type: ignore[attr-defined]
+    convert_checkpoint(config)