Upload model

AbLang_bert_model.py

@@ -0,0 +1,112 @@
+from transformers.models.bert.modeling_bert import BertEncoder, BertPooler, BertEmbeddings, BertForMaskedLM, MaskedLMOutput
+from transformers import BertModel
+from typing import List, Optional, Tuple, Union
+import torch
+
+class BertEmbeddingsV2(BertEmbeddings):
+    def __init__(self, config):
+        super().__init__(config)
+        self.pad_token_id = config.pad_token_id
+        self.position_embeddings = torch.nn.Embedding(config.max_position_embeddings, config.hidden_size, padding_idx=0) # here padding_idx is always 0
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values_length: int = 0,
+    ) -> torch.Tensor:
+        inputs_embeds = self.word_embeddings(input_ids)
+        position_ids = self.create_position_ids_from_input_ids(input_ids)  
+        position_embeddings = self.position_embeddings(position_ids)
+        embeddings = inputs_embeds + position_embeddings
+        return self.dropout(self.LayerNorm(embeddings))
+    
+    def create_position_ids_from_input_ids(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        mask = input_ids.ne(self.pad_token_id).int()
+        return torch.cumsum(mask, dim=1).long() * mask
+    
+
+class BertModelV2(BertModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.embeddings = BertEmbeddingsV2(config)
+
+
+class BertForMaskedLMV2(BertForMaskedLM):
+    def __init__(self, config):
+        super().__init__(config)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
+            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
+        """
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = sequence_output[:, :, 0:24]
+        
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = torch.nn.CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):
+        input_shape = input_ids.shape
+        effective_batch_size = input_shape[0]
+
+        #  add a dummy token
+        if self.config.pad_token_id is None:
+            raise ValueError("The PAD token should be defined for generation")
+
+        attention_mask = torch.cat([attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], dim=-1)
+        dummy_token = torch.full(
+            (effective_batch_size, 1), self.config.pad_token_id, dtype=torch.long, device=input_ids.device
+        )
+        input_ids = torch.cat([input_ids, dummy_token], dim=1)
+
+        return {"input_ids": input_ids, "attention_mask": attention_mask}

config.json

@@ -1,14 +1,13 @@
 {
-  "_name_or_path": "ablang-test",
+  "add_pooling_layer": false,
   "architectures": [
-    "AbLang"
+    "BertModelV2"
   ],
   "attention_probs_dropout_prob": 0.1,
   "auto_map": {
-    "AutoConfig": "config.AbLangConfig",
-    "AutoModel": "model.AbLang"
+    "AutoModel": "AbLang_bert_model.BertModelV2"
   },
-  "chain": "heavy",
+  "classifier_dropout": null,
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
   "hidden_size": 768,
@@ -19,8 +18,11 @@
   "model_type": "bert",
   "num_attention_heads": 12,
   "num_hidden_layers": 12,
-  "ptid": 21,
+  "pad_token_id": 21,
+  "position_embedding_type": "absolute",
   "torch_dtype": "float32",
-  "transformers_version": "4.26.1",
+  "transformers_version": "4.28.1",
+  "type_vocab_size": 2,
+  "use_cache": true,
   "vocab_size": 24
 }

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c526e50a72c294a36b8391ad52936a28832194061026cb1fe9e0455b07e7b5b6
-size 340855773
+oid sha256:c9d5458446b8f723995df81e9b24b7a4635285fcb33d0d787a7e308bb16c75ea
+size 343223341