Support for latest transformers

backbone/config.json

@@ -0,0 +1,27 @@
+{
+  "_name_or_path": "old_models/paraphrase-distilroberta-base-v1/0_Transformer",
+  "architectures": [
+    "RobertaModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 514,
+  "model_type": "roberta",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "output_hidden_states": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.7.0",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 50265
+}

config.json

@@ -13,5 +13,32 @@
   "q_bucket_size": 512,
   "torch_dtype": "float32",
   "transformers_version": "4.33.2",
-  "use_memory_efficient_attention": false
-}
+  "use_memory_efficient_attention": false,
+  "backbone_config": {
+    "_name_or_path": "old_models/paraphrase-distilroberta-base-v1/0_Transformer",
+    "architectures": [
+      "RobertaModel"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "bos_token_id": 0,
+    "eos_token_id": 2,
+    "gradient_checkpointing": false,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_norm_eps": 1e-05,
+    "max_position_embeddings": 514,
+    "model_type": "roberta",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 6,
+    "output_hidden_states": true,
+    "pad_token_id": 1,
+    "position_embedding_type": "absolute",
+    "transformers_version": "4.7.0",
+    "type_vocab_size": 1,
+    "use_cache": true,
+    "vocab_size": 50265
+  }
+}

model.py

@@ -8,7 +8,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange, reduce, repeat
 from torch.utils.checkpoint import checkpoint
-from transformers import AutoModel, PreTrainedModel
+import os
+from transformers import AutoConfig, AutoModel, PreTrainedModel
 
 from .config import LUARConfig
 
@@ -137,9 +138,11 @@ class LUAR(PreTrainedModel):
     """Defines the LUAR model.
     """
     config_class = LUARConfig
-    
+    all_tied_weights_keys = {}
+
     def __init__(self, config):
         super().__init__(config)
+        self._luar_backbone_cfg = dict(config.backbone_config)
         self.create_transformer(revision=config.upstream_transformer_revision)
         self.attn_fn = SelfAttention(
             config.use_memory_efficient_attention,
@@ -148,11 +151,11 @@ class LUAR(PreTrainedModel):
         )
         self.linear = nn.Linear(self.hidden_size, config.embedding_size)
 
-    def create_transformer(self, revision: Optional[str] = None):
-        """Creates the Transformer backbone.
-        """
-        kwargs = {"revision": revision} if revision else {}
-        self.transformer = AutoModel.from_pretrained("sentence-transformers/paraphrase-distilroberta-base-v1", **kwargs)
+    def create_transformer(self, revision=None):
+        """Creates the Transformer backbone (offline: backbone config embedded in LUAR config)."""
+        from transformers import AutoConfig
+        backbone_cfg = AutoConfig.for_model(**self._luar_backbone_cfg)
+        self.transformer = AutoModel.from_config(backbone_cfg)
         self.hidden_size = self.transformer.config.hidden_size
         self.num_attention_heads = self.transformer.config.num_attention_heads
         self.dim_head = self.hidden_size // self.num_attention_heads