Add model code for FiD

fid_model.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+from transformers import AutoModelForSeq2SeqLM, BartConfig, BartForConditionalGeneration, BertModel, BertPreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+
+class FiDEncoder(nn.Module):
+    def __init__(self, encoder):
+        super().__init__()
+        self.encoder = encoder
+
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        # passage_length = input_ids.size(-1)
+        bsz, total_length = input_ids.shape
+        passage_length = total_length // self.n_passages
+        input_ids = input_ids.view(bsz * self.n_passages, passage_length)
+        attention_mask = attention_mask.view(bsz * self.n_passages, passage_length)
+        encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
+
+        encoder_outputs = BaseModelOutput(
+            last_hidden_state=encoder_outputs[0].view(bsz, self.n_passages * passage_length, -1),
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+        return encoder_outputs
+
+
+class FiD(BartForConditionalGeneration):
+    def __init__(self, model_config: BartConfig):
+        if model_config:
+            super().__init__(model_config)
+        else:
+            super().__init__()
+        self.wrap_encoder(self.model.encoder)
+
+    def load_pretrained_params(self, basemodel_name):
+        basemodel = AutoModelForSeq2SeqLM.from_pretrained(basemodel_name)
+        self.model.encoder.encoder.load_state_dict(basemodel.get_encoder().state_dict())
+        self.model.decoder.load_state_dict(basemodel.get_decoder().state_dict())
+        self.lm_head.load_state_dict(basemodel.lm_head.state_dict())
+        print(f"loaded {basemodel_name} parameters.")
+
+    def wrap_encoder(self, encoder):
+        self.model.encoder = FiDEncoder(encoder)
+
+    @classmethod
+    def from_path(cls, model_config):
+
+        return cls(model_config=model_config)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, labels: torch.Tensor = None, **kwargs):
+        """
+
+        Args:
+            input_ids (torch.Tensor): (batch_size, topk, seq_max_len) shape tensor
+            attention_mask (torch.Tensor): (batch_size, topk, seq_max_len) shape tensor
+            labels (torch.Tensor): (batch_size, seq_max_len) shape summarization Tensor
+
+        Return:
+            {
+                logit (torch.Tensor):
+                    (batch_size, max_seq_len, vocab_size) shape logit
+                loss_fn:
+                    logit과 label간의 loss_fn
+                last_hidden_state (torch.Tensor):
+                    (batch_size, max_seq_len, hidden_dim) shape tensor
+            }
+        """
+        if input_ids is not None:
+            if input_ids.ndim == 3:
+                self.model.encoder.n_passages = input_ids.size(1)
+            input_ids = input_ids.view(input_ids.size(0), -1)
+        if attention_mask is not None:
+            attention_mask = attention_mask.view(attention_mask.size(0), -1)
+
+        return super().forward(input_ids=input_ids, attention_mask=attention_mask, labels=labels, **kwargs)
+
+    def generate(self, input_ids, attention_mask, max_length=256):
+        return super().generate(
+            inputs=input_ids.view(input_ids.size(0), -1),
+            attention_mask=attention_mask.view(attention_mask.size(0), -1),
+            max_length=max_length,
+        )