Upload BacformerForMaskedGM

config.json

@@ -1,10 +1,12 @@
 {
-  "_name_or_path": "/rds/user/mw896/rds-flotolab-9X9gY1OFt4M/projects/bacformer/output-data/all-genomes/runs-mgm/12L-full-dataset-rotary-lr15e-5-ampere-new-masking/checkpoint-240000",
   "alpha_contrastive_loss": 0.5,
   "architectures": [
     "BacformerForMaskedGM"
   ],
   "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "configuration_bacformer.BacformerConfig"
+  },
   "batch_size": 1,
   "ckpt_path": null,
   "dataloader_num_workers": 16,
@@ -15,212 +17,14 @@
   "hidden_dropout_prob": 0.1,
   "hidden_size": 480,
   "id2label": {
-    "0": "LABEL_0",
-    "1": "LABEL_1",
-    "2": "LABEL_2",
-    "3": "LABEL_3",
-    "4": "LABEL_4",
-    "5": "LABEL_5",
-    "6": "LABEL_6",
-    "7": "LABEL_7",
-    "8": "LABEL_8",
-    "9": "LABEL_9",
-    "10": "LABEL_10",
-    "11": "LABEL_11",
-    "12": "LABEL_12",
-    "13": "LABEL_13",
-    "14": "LABEL_14",
-    "15": "LABEL_15",
-    "16": "LABEL_16",
-    "17": "LABEL_17",
-    "18": "LABEL_18",
-    "19": "LABEL_19",
-    "20": "LABEL_20",
-    "21": "LABEL_21",
-    "22": "LABEL_22",
-    "23": "LABEL_23",
-    "24": "LABEL_24",
-    "25": "LABEL_25",
-    "26": "LABEL_26",
-    "27": "LABEL_27",
-    "28": "LABEL_28",
-    "29": "LABEL_29",
-    "30": "LABEL_30",
-    "31": "LABEL_31",
-    "32": "LABEL_32",
-    "33": "LABEL_33",
-    "34": "LABEL_34",
-    "35": "LABEL_35",
-    "36": "LABEL_36",
-    "37": "LABEL_37",
-    "38": "LABEL_38",
-    "39": "LABEL_39",
-    "40": "LABEL_40",
-    "41": "LABEL_41",
-    "42": "LABEL_42",
-    "43": "LABEL_43",
-    "44": "LABEL_44",
-    "45": "LABEL_45",
-    "46": "LABEL_46",
-    "47": "LABEL_47",
-    "48": "LABEL_48",
-    "49": "LABEL_49",
-    "50": "LABEL_50",
-    "51": "LABEL_51",
-    "52": "LABEL_52",
-    "53": "LABEL_53",
-    "54": "LABEL_54",
-    "55": "LABEL_55",
-    "56": "LABEL_56",
-    "57": "LABEL_57",
-    "58": "LABEL_58",
-    "59": "LABEL_59",
-    "60": "LABEL_60",
-    "61": "LABEL_61",
-    "62": "LABEL_62",
-    "63": "LABEL_63",
-    "64": "LABEL_64",
-    "65": "LABEL_65",
-    "66": "LABEL_66",
-    "67": "LABEL_67",
-    "68": "LABEL_68",
-    "69": "LABEL_69",
-    "70": "LABEL_70",
-    "71": "LABEL_71",
-    "72": "LABEL_72",
-    "73": "LABEL_73",
-    "74": "LABEL_74",
-    "75": "LABEL_75",
-    "76": "LABEL_76",
-    "77": "LABEL_77",
-    "78": "LABEL_78",
-    "79": "LABEL_79",
-    "80": "LABEL_80",
-    "81": "LABEL_81",
-    "82": "LABEL_82",
-    "83": "LABEL_83",
-    "84": "LABEL_84",
-    "85": "LABEL_85",
-    "86": "LABEL_86",
-    "87": "LABEL_87",
-    "88": "LABEL_88",
-    "89": "LABEL_89",
-    "90": "LABEL_90",
-    "91": "LABEL_91",
-    "92": "LABEL_92",
-    "93": "LABEL_93",
-    "94": "LABEL_94",
-    "95": "LABEL_95",
-    "96": "LABEL_96",
-    "97": "LABEL_97",
-    "98": "LABEL_98",
-    "99": "LABEL_99"
+    "0": "LABEL_0"
   },
   "initializer_range": 0.02,
   "input_dir": "/rds/user/mw896/rds-flotolab-9X9gY1OFt4M/projects/bacformer/input-data/eval-genomes/",
   "intermediate_size": 1280,
   "is_causal_gm": false,
   "label2id": {
-    "LABEL_0": 0,
-    "LABEL_1": 1,
-    "LABEL_10": 10,
-    "LABEL_11": 11,
-    "LABEL_12": 12,
-    "LABEL_13": 13,
-    "LABEL_14": 14,
-    "LABEL_15": 15,
-    "LABEL_16": 16,
-    "LABEL_17": 17,
-    "LABEL_18": 18,
-    "LABEL_19": 19,
-    "LABEL_2": 2,
-    "LABEL_20": 20,
-    "LABEL_21": 21,
-    "LABEL_22": 22,
-    "LABEL_23": 23,
-    "LABEL_24": 24,
-    "LABEL_25": 25,
-    "LABEL_26": 26,
-    "LABEL_27": 27,
-    "LABEL_28": 28,
-    "LABEL_29": 29,
-    "LABEL_3": 3,
-    "LABEL_30": 30,
-    "LABEL_31": 31,
-    "LABEL_32": 32,
-    "LABEL_33": 33,
-    "LABEL_34": 34,
-    "LABEL_35": 35,
-    "LABEL_36": 36,
-    "LABEL_37": 37,
-    "LABEL_38": 38,
-    "LABEL_39": 39,
-    "LABEL_4": 4,
-    "LABEL_40": 40,
-    "LABEL_41": 41,
-    "LABEL_42": 42,
-    "LABEL_43": 43,
-    "LABEL_44": 44,
-    "LABEL_45": 45,
-    "LABEL_46": 46,
-    "LABEL_47": 47,
-    "LABEL_48": 48,
-    "LABEL_49": 49,
-    "LABEL_5": 5,
-    "LABEL_50": 50,
-    "LABEL_51": 51,
-    "LABEL_52": 52,
-    "LABEL_53": 53,
-    "LABEL_54": 54,
-    "LABEL_55": 55,
-    "LABEL_56": 56,
-    "LABEL_57": 57,
-    "LABEL_58": 58,
-    "LABEL_59": 59,
-    "LABEL_6": 6,
-    "LABEL_60": 60,
-    "LABEL_61": 61,
-    "LABEL_62": 62,
-    "LABEL_63": 63,
-    "LABEL_64": 64,
-    "LABEL_65": 65,
-    "LABEL_66": 66,
-    "LABEL_67": 67,
-    "LABEL_68": 68,
-    "LABEL_69": 69,
-    "LABEL_7": 7,
-    "LABEL_70": 70,
-    "LABEL_71": 71,
-    "LABEL_72": 72,
-    "LABEL_73": 73,
-    "LABEL_74": 74,
-    "LABEL_75": 75,
-    "LABEL_76": 76,
-    "LABEL_77": 77,
-    "LABEL_78": 78,
-    "LABEL_79": 79,
-    "LABEL_8": 8,
-    "LABEL_80": 80,
-    "LABEL_81": 81,
-    "LABEL_82": 82,
-    "LABEL_83": 83,
-    "LABEL_84": 84,
-    "LABEL_85": 85,
-    "LABEL_86": 86,
-    "LABEL_87": 87,
-    "LABEL_88": 88,
-    "LABEL_89": 89,
-    "LABEL_9": 9,
-    "LABEL_90": 90,
-    "LABEL_91": 91,
-    "LABEL_92": 92,
-    "LABEL_93": 93,
-    "LABEL_94": 94,
-    "LABEL_95": 95,
-    "LABEL_96": 96,
-    "LABEL_97": 97,
-    "LABEL_98": 98,
-    "LABEL_99": 99
+    "LABEL_0": 0
   },
   "layer_norm_eps": 1e-12,
   "logging_steps": 500,
@@ -263,7 +67,7 @@
   "test_after_train": false,
   "torch_dtype": "float32",
   "train_subset_prop": 1.0,
-  "transformers_version": "4.38.2",
+  "transformers_version": "4.50.3",
   "warmup_proportion": 0.1,
   "weight_decay": 0.01
 }

configuration_bacformer.py

@@ -0,0 +1,72 @@
+from typing import Literal
+
+from transformers import PretrainedConfig
+
+SPECIAL_TOKENS_DICT = {
+    "PAD": 0,
+    "MASK": 1,
+    "CLS": 2,
+    "SEP": 3,
+    "PROT_EMB": 4,
+    "END": 5,
+}
+
+
+class BacformerConfig(PretrainedConfig):
+    """Configuration class to store the configuration of a `BacformerModel`."""
+
+    model_type = "bacformer"
+
+    def __init__(
+        self,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 8,
+        hidden_size: int = 480,  # default esm2_t12_35M_UR50D embedding dim
+        intermediate_size: int = 1280,
+        hidden_dropout_prob: float = 0.1,
+        attention_probs_dropout_prob: float = 0.1,
+        max_position_embeddings: int = 6000,
+        max_token_type_embeddings: int = 1000,
+        layer_norm_eps: float = 1e-12,
+        initializer_range: float = 0.02,
+        pad_token_id: int = SPECIAL_TOKENS_DICT["PAD"],
+        mask_token_id: int = SPECIAL_TOKENS_DICT["MASK"],
+        prot_emb_token_id: int = SPECIAL_TOKENS_DICT["PROT_EMB"],
+        end_token_id: int = SPECIAL_TOKENS_DICT["END"],
+        num_special_tokens: int = len(SPECIAL_TOKENS_DICT),
+        protein_clusters_vocab_size: int = 50001,  # equal to the nr of protein clusters + 1
+        num_labels: int = 1,  # for downstream tasks
+        is_causal_gm: bool = False,
+        return_dict: bool = False,
+        return_attn_weights: bool = False,
+        alpha_contrastive_loss: float = 0.5,
+        # only to use in the BacformerForGenomeClassification
+        problem_type: Literal[
+            "regression", "binary_classification", "single_label_classification", "multi_label_classification"
+        ] = "single_label_classification",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.max_token_type_embeddings = max_token_type_embeddings
+        self.layer_norm_eps = layer_norm_eps
+        self.initializer_range = initializer_range
+        self.pad_token_id = pad_token_id
+        self.mask_token_id = mask_token_id
+        self.prot_emb_token_id = prot_emb_token_id
+        self.end_token_id = end_token_id
+        self.num_special_tokens = num_special_tokens
+        self.protein_clusters_vocab_size = protein_clusters_vocab_size
+        self.num_labels = num_labels
+        self.is_causal_gm = is_causal_gm
+        self.return_dict = return_dict
+        self.return_attn_weights = return_attn_weights
+        self.problem_type = problem_type
+        self.alpha_contrastive_loss = alpha_contrastive_loss