Upload P5Pretraining

config.json

@@ -0,0 +1,67 @@
+{
+  "_name_or_path": "t5-small",
+  "activation_dropout": 0.1,
+  "architectures": [
+    "P5Pretraining"
+  ],
+  "attention_dropout": 0.1,
+  "auto_map": {
+    "AutoModel": "pretrain_model.P5Pretraining"
+  },
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 512,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "relu",
+  "dropout": 0.1,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "relu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": false,
+  "layer_norm_epsilon": 1e-06,
+  "losses": "rating,sequential,explanation,review,traditional",
+  "model_type": "t5",
+  "n_positions": 512,
+  "num_decoder_layers": 6,
+  "num_heads": 8,
+  "num_layers": 6,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "task_specific_params": {
+    "summarization": {
+      "early_stopping": true,
+      "length_penalty": 2.0,
+      "max_length": 200,
+      "min_length": 30,
+      "no_repeat_ngram_size": 3,
+      "num_beams": 4,
+      "prefix": "summarize: "
+    },
+    "translation_en_to_de": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to German: "
+    },
+    "translation_en_to_fr": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to French: "
+    },
+    "translation_en_to_ro": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to Romanian: "
+    }
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.25.1",
+  "use_cache": true,
+  "vocab_size": 32100
+}

pretrain_model.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from modeling_p5 import P5
+
+class P5Pretraining(P5):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.losses = self.config.losses.split(',')
+
+    def train_step(self, batch):
+
+        device = next(self.parameters()).device
+        input_ids = batch['input_ids'].to(device)
+        whole_word_ids = batch['whole_word_ids'].to(device)
+
+        lm_labels = batch["target_ids"].to(device)
+
+        loss_weights = batch["loss_weights"].to(device)
+
+        output = self(
+            input_ids=input_ids,
+            whole_word_ids=whole_word_ids,
+            labels=lm_labels,
+            return_dict=True
+        )
+        assert 'loss' in output
+
+        lm_mask = lm_labels != -100
+        lm_mask = lm_mask.float()
+        B, L = lm_labels.size()
+
+        loss = output['loss']
+
+        loss = loss.view(B, L) * lm_mask
+
+        loss = loss.sum(dim=1) / lm_mask.sum(dim=1).clamp(min=1)
+
+        task_counts = {task: 0 for task in self.losses}
+        task_loss = {task: 0 for task in self.losses}
+
+        results = {}
+
+        results['loss'] = (loss * loss_weights).mean()
+        results['total_loss'] = loss.detach().sum()
+        results['total_loss_count'] = len(loss)
+
+        task_counts = {task: 0 for task in self.losses}
+        task_loss = {task: 0 for task in self.losses}
+
+        for _loss, task in zip(loss.detach(), batch['task']):
+            task_loss[task] += _loss
+            task_counts[task] += 1
+
+        for task in self.losses:
+            if task_counts[task] > 0:
+                results[f'{task}_loss'] = task_loss[task]
+                results[f'{task}_loss_count'] = task_counts[task]
+
+        return results
+
+    @torch.no_grad()
+    def valid_step(self, batch):
+        self.eval()
+        device = next(self.parameters()).device
+        input_ids = batch['input_ids'].to(device)
+
+        lm_labels = batch["target_ids"].to(device)
+
+        loss_weights = batch["loss_weights"].to(device)
+
+        output = self(
+            input_ids=input_ids,
+            labels=lm_labels,
+            return_dict=True
+        )
+        assert 'loss' in output
+
+        lm_mask = lm_labels != -100
+        lm_mask = lm_mask.float()
+        B, L = lm_labels.size()
+
+        loss = output['loss']
+
+        loss = loss.view(B, L) * lm_mask
+
+        loss = loss.sum(dim=1) / lm_mask.sum(dim=1).clamp(min=1)
+
+        results = {}
+
+        results['loss'] = (loss * loss_weights).mean()
+        results['total_loss'] = loss.detach().sum()
+        results['total_loss_count'] = len(loss)
+
+        task_counts = {task: 0 for task in self.losses}
+        task_loss = {task: 0 for task in self.losses}
+
+        for _loss, task in zip(loss.detach(), batch['task']):
+            task_loss[task] += _loss
+            task_counts[task] += 1
+
+        for task in self.losses:
+            if task_counts[task] > 0:
+                results[f'{task}_loss'] = task_loss[task]
+                results[f'{task}_loss_count'] = task_counts[task]
+
+        if 'rating' in self.losses:
+            output = self.generate(
+                input_ids=input_ids
+            )
+
+            generated_score = self.tokenizer.batch_decode(output, skip_special_tokens=True)
+
+            results['rating_pred'] = generated_score
+
+        return results
+
+    @torch.no_grad()
+    def generate_step(self, batch):
+        self.eval()
+        device = next(self.parameters()).device
+        input_ids = batch['input_ids'].to(device)
+
+        output = self.generate(
+            input_ids=input_ids,
+        )
+
+        generated_sents = self.tokenizer.batch_decode(output, skip_special_tokens=True)
+
+        return generated_sents

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1077d7f9d23ba28012cef084386d83491f9fe01da05089780da3530ac74123d8
+size 243063197