Delete pretrain_model.py

pretrain_model.py

@@ -1,133 +0,0 @@
-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