CodeSumBART-ForGeneration.py

""" BART (Seq-2-Seq) Encoder/Decoder Transformer using PyTorch Lightning.
    We tuned this model on the Funcom Dataset.
    From template code found at: https://lightning-transformers.readthedocs.io
                                 /en/latest/tasks/nlp/summarization.html
                            and: https://gist.github.com/ajazturki10
                                 /247ac21e001025d8b65c7418edd4faf5#file-model-py

Author: Jesse Phillips <j.m.phillips@lancaster.ac.uk>
"""
import json
import os
import torch
import time
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from lightning_transformers.task.nlp.summarization import (
    SummarizationDataModule)
from transformers import BartForConditionalGeneration, \
                         AutoTokenizer, \
                         AutoConfig
import evaluate
import gc
from meteor import Meteor

base = "facebook/bart-base"
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=base)
ignoreID = torch.tensor(-100)
padTokenID = torch.tensor(tokenizer.pad_token_id)


class AdamWBARTBleuModel(pl.LightningModule):
    def getMetric(self, metricName: str) -> evaluate.EvaluationModule:
        """ Gets a metric from HuggingFace's Evaluate API.
            Tries three times because their network can get flaky when busy.
        Args:
            metricName (str): the name of the metric to use.
        Returns:
            EvaluationModule: the metric.
        """
        try:
            return evaluate.load(metricName)
        except Exception:
            time.sleep(60)
            try:
                return evaluate.load(metricName)
            except Exception:
                time.sleep(60)
                try:
                    return evaluate.load(metricName)
                except Exception as e:
                    print(f"could not access HuggingFace {metricName}")
                    raise e

    def __init__(self, model: str, *args: tuple, **kwargs: dict) -> None:
        """ model constructor.
        Args:
            model (str): the base name for the HuggingFace model.
            args (tuple): model args.
            kwargs (dict): model keyword args.
        """
        super(AdamWBARTBleuModel, self).__init__()
        self.use_stemmer = kwargs["use_stemmer"]
        self.val_target_max_length = kwargs["val_target_max_length"]
        self.num_beams = kwargs["num_beams"]
        self.compute_generate_metrics = kwargs["compute_generate_metrics"]

        config = AutoConfig.from_pretrained(model)
        self.model = BartForConditionalGeneration(config=config)
        self.model.init_weights()
        self.save_hyperparameters(logger=False)

    def configure_optimizers(self) -> torch.optim.AdamW:
        """ configures the optimizer(s) for our model (AdamW)
        Returns:
            AdamW: the optimizer used.
        """
        return torch.optim.AdamW(self.parameters(), lr=2e-5)

    def forward(self,
                input_ids: torch.Tensor,
                attention_mask: torch.Tensor,
                labels=None,
                decoder_attention_mask=None) -> tuple:
        """ forward step for our model
        Args:
            input_ids (torch.Tensor): the inputs
            attention_mask (torch.Tensor): the attention mask
            labels (None/torch.Tensor): the target labels
            decoder_attention_mask (None/torch.tensor): dec. att. mask if used
        Returns:
            Tuple: the loss and logits.
        """
        outputs = self.model(input_ids=input_ids,
                             attention_mask=attention_mask,
                             labels=labels,
                             decoder_attention_mask=decoder_attention_mask)
        return outputs.loss, outputs.logits

    def on_train_epoch_start(self):
        if os.path.exists(checkpointCallback.best_model_path):
            oldWeights = torch.load(checkpointCallback.best_model_path
                                    )['state_dict']
            model.load_state_dict(oldWeights)
            self.model.train()

    def training_step(self, batch: dict, batch_idx: int) -> torch.Tensor:
        """ The training loop for our model.
        NB: call self or forward, not both.
        Args:
            batch (dict): the data batch
            batch_idx (int): the batch index
        Returns:
            torch.Tensor: the loss.
        """
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["labels"]

        loss, output = self(input_ids, attention_mask, labels)
        return loss

    def validation_step(self, batch: dict, batch_idx: int) -> float:
        """ The validation loop for our model.
        Args:
            batch (dict): the data batch
            batch_idx (int): the batch index
        Returns:
            float: the BLEU-4 score.
        """
        labels = batch["labels"]
        labels = labels.where(labels != ignoreID.to(self.device),
                              padTokenID.to(self.device))
        decodedLabels = tokenizer.batch_decode(labels,
                                               skip_special_tokens=True)
        decodedLabels = [[label.strip()] for label in decodedLabels]
        batch = {k: torch.Tensor(v) for k, v in batch.items()}
        outputs = self.model(**batch)
        logits = outputs.logits
        decodedPreds = tokenizer.batch_decode(torch.argmax(logits, dim=-1),
                                              skip_special_tokens=True)
        decodedPreds = [pred.strip() for pred in decodedPreds]
        metric = self.getMetric("bleu")
        bleu = metric.compute(predictions=decodedPreds,
                              references=decodedLabels,
                              max_order=4,
                              smooth=False)["bleu"]
        self.log("val_bleu", bleu, on_step=False, on_epoch=True, prog_bar=True)
        return bleu

    def test_step(self, batch: dict, batch_idx: int) -> torch.Tensor:
        """ Post-training evaluation.
        Args:
            batch (dict): the data batch
            batch_idx (int): the batch index
        Returns:
            dict: the outputs.
        """
        torch.cuda.empty_cache()
        gc.collect()
        with torch.no_grad():
            output = self.model(**batch)
        logits = output.logits
        predictions = tokenizer.batch_decode(torch.argmax(logits, dim=-1),
                                             skip_special_tokens=True)
        predictions = [pred.strip() for pred in predictions]

        labels = batch["labels"]
        labels = labels.where(labels != ignoreID.to(self.device),
                              padTokenID.to(self.device))
        references = tokenizer.batch_decode(labels,
                                            skip_special_tokens=True)
        references = [label.strip() for label in references]
        with open("/mnt/nfs/homes/phillijm/dataset/evaluation/dataset.json",
                  "r",
                  encoding="UTF-8") as fp1:
            data = json.load(fp1)

        with open("/mnt/nfs/homes/phillijm/LLM_Sum/Inputs/outputs.txt",
                  "a",
                  encoding="UTF-8") as fp:
            for cnt in range(len(predictions)):
                cleanPred = predictions[cnt].replace("\n", " ")
                cleanRef = references[cnt].replace("\n", " ")
                code = ""
                for pair in data:
                    if (pair["target"] == references[cnt] + '\n'):
                        code = pair["source"]
                        break
                if code != "":
                    code = code.replace("\n", " ")
                    fp.write(f"ITEM:\nmethod: {code}\nprediction: {cleanPred}\nreference: {cleanRef}\n\n")
        return

    def on_test_epoch_end(self):
        pass


class FuncomDataModule(SummarizationDataModule):
    def __init__(self, *args: tuple, **kwargs: dict) -> None:
        super(FuncomDataModule, self).__init__(*args, **kwargs)


if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"  # GPUs we're allowed to use.
    torch.multiprocessing.freeze_support()

    torch.cuda.empty_cache()  # Clear the PyTorch cache - saves GPU memory.
    torch.backends.cudnn.benchmark = True  # Run fastest convolutions.
    gc.collect()

    trainingDataPath = "/mnt/nfs/homes/phillijm/dataset/training/dataset.json"
    validationDataPath = "/mnt/nfs/homes/phillijm/dataset/validation/dataset.json"
    evaluationDataPath = "/mnt/nfs/homes/phillijm/dataset/evaluation/dataset.json"

    checkpointCallback = ModelCheckpoint(monitor="val_bleu",
                                         save_top_k=3,
                                         mode="max",
                                         save_weights_only=True)

    dm = FuncomDataModule(batch_size=48,
                          num_workers=4,
                          max_source_length=128,
                          max_target_length=128,
                          train_file=os.path.abspath(trainingDataPath),
                          validation_file=os.path.abspath(validationDataPath),
                          test_file=os.path.abspath(evaluationDataPath),
                          tokenizer=tokenizer,
                          max_length=512,
                          padding="max_length")

    model = AdamWBARTBleuModel(model=base,
                               use_stemmer=True,
                               val_target_max_length=128,
                               num_beams=None,
                               compute_generate_metrics=True)

    trainer = pl.Trainer(accelerator="gpu",
                         devices=[0],  # GPU ID/s GO HERE
                         check_val_every_n_epoch=1,
                         max_epochs=0,
                         min_epochs=0,
                         callbacks=[checkpointCallback,
                                    EarlyStopping(monitor="val_bleu",
                                                  mode="max",
                                                  patience=5)],
                         limit_train_batches=0,
                         limit_val_batches=0)
    trainer.fit(model, dm)

    # Evaluate the model.
    trainer.test(ckpt_path="full_csb.ckpt", datamodule=dm)