Upload RankingPrompterForPreTraining

config.json

@@ -1,7 +1,7 @@
 {
   "_name_or_path": "D://huggingface_model/RankingPrompterForPreTraining-small",
   "architectures": [
-    "UMT5Model"
+    "RankingPrompterForPreTraining"
   ],
   "auto_map": {
     "AutoConfig": "configuration_rankingprompter.RankingPrompterConfig",

configuration_rankingprompter.py

@@ -0,0 +1,82 @@
+from transformers import PretrainedConfig
+
+class RankingPrompterConfig(PretrainedConfig):
+    model_type = "umt5"
+
+    def __init__(
+        self,
+        vocab_size=250112,
+        d_model=512,
+        d_kv=64,
+        d_ff=1024,
+        num_layers=8,
+        num_decoder_layers=None,
+        num_heads=6,
+        relative_attention_num_buckets=32,
+        relative_attention_max_distance=128,
+        dropout_rate=0.1,
+        layer_norm_epsilon=1e-6,
+        initializer_factor=1.0,
+        feed_forward_proj="gated-gelu",
+        is_encoder_decoder=True,
+        use_cache=True,
+        tokenizer_class="T5Tokenizer",
+        tie_word_embeddings=True,
+        pad_token_id=0,
+        eos_token_id=1,
+        decoder_start_token_id=0,
+        classifier_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(
+            is_encoder_decoder=is_encoder_decoder,
+            tokenizer_class=tokenizer_class,
+            tie_word_embeddings=tie_word_embeddings,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.d_kv = d_kv
+        self.d_ff = d_ff
+        self.num_layers = num_layers
+        self.num_decoder_layers = (
+            num_decoder_layers if num_decoder_layers is not None else self.num_layers
+        )  # default = symmetry
+        self.num_heads = num_heads
+        self.relative_attention_num_buckets = relative_attention_num_buckets
+        self.relative_attention_max_distance = relative_attention_max_distance
+        self.dropout_rate = dropout_rate
+        self.classifier_dropout = classifier_dropout
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_factor = initializer_factor
+        self.feed_forward_proj = feed_forward_proj
+        self.use_cache = use_cache
+
+        act_info = self.feed_forward_proj.split("-")
+        self.dense_act_fn = act_info[-1]
+        self.is_gated_act = act_info[0] == "gated"
+
+        if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2:
+            raise ValueError(
+                f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
+                "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
+                "'gated-gelu' or 'relu'"
+            )
+
+        if feed_forward_proj == "gated-gelu":
+            self.dense_act_fn = "gelu_new"
+
+    @property
+    def hidden_size(self):
+        return self.d_model
+
+    @property
+    def num_attention_heads(self):
+        return self.num_heads
+
+    @property
+    def num_hidden_layers(self):
+        return self.num_layers

modeling_rankingprompter.py

@@ -0,0 +1,140 @@
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import UMT5Model
+from .configuration_rankingprompter import RankingPrompterConfig
+
+
+@dataclass
+class RankingPrompterForPreTrainingOutput:
+    loss: torch.FloatTensor = None
+    logits: torch.FloatTensor = None
+
+
+class RankingPrompterForPreTraining(UMT5Model):
+    config_class = RankingPrompterConfig
+
+    _tied_weights_keys = [
+        "encoder.embed_tokens.weight",
+        "decoder.embed_tokens.weight",
+    ]
+
+    def __init__(self, config):
+        # encoder, decoder and shared are from UMT5Model
+        super().__init__(config)
+
+        # add ranking head
+        self.ranking_head = nn.Linear(config.d_model, 1)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        # ctx for mixed precision training
+        self.ctx = nullcontext()
+
+    def enable_amp_ctx(self, device_type="cuda", dtype=torch.bfloat16):
+        self.ctx = torch.amp.autocast(device_type=device_type, dtype=dtype)
+
+    def disable_amp_ctx(self):
+        self.ctx = nullcontext()
+
+    def forward(
+        self,
+        document_input_ids: Optional[torch.LongTensor] = None,
+        document_attention_mask: Optional[torch.FloatTensor] = None,
+        question_input_ids: Optional[torch.LongTensor] = None,
+        question_attention_mask: Optional[torch.BoolTensor] = None,
+        encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.FloatTensor], RankingPrompterForPreTrainingOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ...,
+            config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for
+            labels in `[0, ..., config.vocab_size]`
+
+        Returns:
+
+        ```"""
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        # document_input_ids: [batch_size, num_doc, doc_seq_len]
+        batch_size, num_doc, doc_seq_len = document_input_ids.shape
+        #
+        document_input_ids = document_input_ids.view(-1, doc_seq_len)
+        # to [batch_size * num_doc, doc_seq_len]
+        document_attention_mask = document_attention_mask.view(-1, doc_seq_len)
+
+        # Convert encoder inputs in embeddings if needed
+        with self.ctx:
+            encoder_outputs = self.encoder(
+                input_ids=document_input_ids,
+                attention_mask=document_attention_mask,
+                return_dict=return_dict,
+            )
+
+        document_embeds = encoder_outputs[0]
+
+        # repeat question inputs for each document
+        # question_input_ids: [batch_size, question_seq_len]
+        question_seq_len = question_input_ids.shape[1]
+        question_input_ids = (
+            question_input_ids.unsqueeze(1)
+            .expand(-1, num_doc, -1)
+            .reshape(-1, question_seq_len)
+        )  # [batch_size * num_doc, question_seq_len]
+        question_attention_mask = (
+            question_attention_mask.unsqueeze(1)
+            .expand(-1, num_doc, -1)
+            .reshape(-1, question_seq_len)
+        )  # [batch_size * num_doc, question_seq_len]
+
+        # Decode
+        with self.ctx:
+            decoder_outputs = self.decoder(
+                input_ids=question_input_ids,
+                attention_mask=question_attention_mask,
+                past_key_values=past_key_values,
+                encoder_hidden_states=document_embeds,
+                encoder_attention_mask=document_attention_mask,
+                use_cache=use_cache,
+                return_dict=return_dict,
+            )
+        # [batch_size * num_doc, soft_prompt_len + question_seq_len, hidden_size]
+        sequence_output = decoder_outputs[0]
+        # [batch_size * num_doc, soft_prompt_len, hidden_size]
+        question_seq_len = sequence_output.size(1)
+        # [batch_size, num_doc, soft_prompt_len, hidden_size]
+        soft_prompt_output = sequence_output.view(
+            batch_size, num_doc, question_seq_len, -1
+        )
+
+        # [batch_size, num_doc, self.num_soft_prompt_tokens, hidden_size] -> [batch_size, num_doc, hidden_size]
+        ranking_logits = self.ranking_head(soft_prompt_output.mean(dim=2))
+
+        # rank loss
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss(ignore_index=-100)
+            ranking_logits = ranking_logits.view(batch_size, num_doc)
+            loss = loss_fct(ranking_logits, labels)
+
+        if not return_dict:
+            output = (ranking_logits,) + decoder_outputs[1:] + encoder_outputs
+            return ((loss,) + output) if loss is not None else output
+
+        return RankingPrompterForPreTrainingOutput(
+            loss=loss,
+            logits=ranking_logits
+        )
+
+

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:982317cd72cc9fc369542b1145e9141fc19168206fee0dbdcf91fc3d9ff0c2e0
-size 701400953
+oid sha256:b90ef8ceeeffc7b033e65dfc28f3adf8d82cbdad204df0677ae0c0f45f4f0c24
+size 701403585