Support specifying upstream sentence transformers model revision

config.json

@@ -8,6 +8,7 @@
   },
   "embedding_size": 512,
   "k_bucket_size": 1024,
+  "upstream_transformer_revision": null,
   "model_type": "LUAR",
   "q_bucket_size": 512,
   "torch_dtype": "float32",

config.py

@@ -1,4 +1,6 @@
 
+from typing import Optional
+
 from transformers import PretrainedConfig
 
 class LUARConfig(PretrainedConfig):
@@ -9,10 +11,12 @@ class LUARConfig(PretrainedConfig):
         use_memory_efficient_attention=False,
         q_bucket_size=512,
         k_bucket_size=1024,
+        upstream_transformer_revision: Optional[str] = None,
         **kwargs,
     ):
         self.embedding_size = embedding_size
         self.use_memory_efficient_attention = use_memory_efficient_attention
         self.q_bucket_size = q_bucket_size
         self.k_bucket_size = k_bucket_size
-        super().__init__(**kwargs)
+        self.upstream_transformer_revision = upstream_transformer_revision
+        super().__init__(**kwargs)

model.py

@@ -1,6 +1,7 @@
 
 import math
 from functools import partial
+from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -139,7 +140,7 @@ class LUAR(PreTrainedModel):
     
     def __init__(self, config):
         super().__init__(config)
-        self.create_transformer()
+        self.create_transformer(revision=config.upstream_transformer_revision)
         self.attn_fn = SelfAttention(
             config.use_memory_efficient_attention,
             config.q_bucket_size,
@@ -147,10 +148,11 @@ class LUAR(PreTrainedModel):
         )
         self.linear = nn.Linear(self.hidden_size, config.embedding_size)
 
-    def create_transformer(self):
+    def create_transformer(self, revision: Optional[str] = None):
         """Creates the Transformer backbone.
         """
-        self.transformer = AutoModel.from_pretrained("sentence-transformers/paraphrase-distilroberta-base-v1")
+        kwargs = {"revision": revision} if revision else {}
+        self.transformer = AutoModel.from_pretrained("sentence-transformers/paraphrase-distilroberta-base-v1", **kwargs)
         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