add model

config.json

@@ -1,23 +1,22 @@
 {
+  "_name_or_path": "nz/RITA_xl",
   "architectures": [
     "RITAModel"
   ],
   "auto_map": {
     "AutoConfig": "rita_configuration.RITAConfig",
-    "AutoModel": "rita_modeling.RITAModel"
+    "AutoModel": "rita_modeling.RITAModel",
+    "AutoModelForCausalLM": "rita_modeling.RITAModel"
   },
-  "bos_token_id": [
-    50256
-  ],
   "d_feedforward": 8192,
   "d_model": 2048,
   "dropout": 0.0,
   "eos_token_id": 50256,
   "max_seq_len": 1024,
-  "model_type": "codegen",
+  "model_type": "rita",
   "num_heads": 32,
   "num_layers": 24,
-  "torch_dtype": "float32",
+  "torch_dtype": "float16",
   "transformers_version": "4.18.0",
-  "vocab_size": 128
+  "vocab_size": 26
 }

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bd3eab5b0d211f648b4dc7ab99291186c6866ff533ebdd94b1ddd0334c118f5d
-size 4836636593
+oid sha256:295ec5633c129e34d2a83ffb6adced9dcdb6cdae3f42534c85b0e4ed7adbfb21
+size 2417438283

rita_configuration.py

@@ -1,26 +1,24 @@
-
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
 
 class RITAConfig(PretrainedConfig):
-    model_type = "codegen"
+    model_type = "rita"
 
     def __init__(
         self,
-        vocab_size=128,
+        vocab_size=26,
         d_model=768,
         num_layers=12,
         max_seq_len=1024,
         num_heads=12,
         dropout=0.,
         ff_ratio=4,
-        bos_token_id=50256, # TODO
-        eos_token_id=50256, # TODO
+        eos_token_id=2,
         **kwargs,
     ):
-        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+        super().__init__(eos_token_id=eos_token_id, **kwargs)
         self.vocab_size = vocab_size
         self.d_model = d_model
         self.num_heads = num_heads
@@ -28,5 +26,4 @@ class RITAConfig(PretrainedConfig):
         self.num_layers = num_layers
         self.max_seq_len=max_seq_len
         self.dropout = dropout
-        self.bos_token_id=bos_token_id,
-        self.eos_token_id=eos_token_id
+        self.eos_token_id=eos_token_id

rita_modeling.py

@@ -222,10 +222,10 @@ class RITAModel(PreTrainedModel):
         self.final_norm = nn.LayerNorm(config.d_model)
         self.projector = nn.Linear(config.d_model, config.vocab_size, bias = False)
 
-    def forward(self, ids, attn_mask=None, padding_mask=None, return_hidden=False) -> torch.FloatTensor:
-        x = self.embedding(ids)  # N x L x D
+    def forward(self, input_ids, attn_mask=None, padding_mask=None, return_hidden=False) -> torch.FloatTensor:
+        x = self.embedding(input_ids)  # N x L x D
         if attn_mask == None:
-            attn_mask = (torch.triu(torch.ones(ids.size(1), ids.size(1))) == 0).transpose(0, 1).contiguous()
+            attn_mask = (torch.triu(torch.ones(input_ids.size(1), input_ids.size(1))) == 0).transpose(0, 1).contiguous().to(input_ids.device)
         for layer in self.layers:
             x = layer(x, attn_mask=attn_mask, padding_mask=padding_mask)
         x = self.final_norm(x)  # N x L x D
@@ -246,5 +246,4 @@ class RITAModel(PreTrainedModel):
         return self.projector
 
     def set_output_embeddings(self, new_projector):
-        return new_projector
-            
+        self.projector = new_projector