Create BaseTrainer.py

Nested/trainers/BaseTrainer.py

@@ -0,0 +1,125 @@
+import os
+import torch
+import logging
+import natsort
+import glob
+from huggingface_hub import hf_hub_download, snapshot_download
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTrainer:
+    def __init__(
+        self,
+        model=None,
+        max_epochs=50,
+        optimizer=None,
+        scheduler=None,
+        loss=None,
+        train_dataloader=None,
+        val_dataloader=None,
+        test_dataloader=None,
+        log_interval=10,
+        summary_writer=None,
+        output_path=None,
+        clip=5,
+        patience=5
+    ):
+        self.model = model
+        self.max_epochs = max_epochs
+        self.train_dataloader = train_dataloader
+        self.val_dataloader = val_dataloader
+        self.test_dataloader = test_dataloader
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.loss = loss
+        self.log_interval = log_interval
+        self.summary_writer = summary_writer
+        self.output_path = output_path
+        self.current_timestep = 0
+        self.current_epoch = 0
+        self.clip = clip
+        self.patience = patience
+
+    def tag(self, dataloader, is_train=True):
+        """
+        Given a dataloader containing segments, predict the tags
+        :param dataloader: torch.utils.data.DataLoader
+        :param is_train: boolean - True for training model, False for evaluation
+        :return: Iterator
+                    subwords (B x T x NUM_LABELS)- torch.Tensor - BERT subword ID
+                    gold_tags (B x T x NUM_LABELS) - torch.Tensor - ground truth tags IDs
+                    tokens - List[Nested.data.dataset.Token] - list of tokens
+                    valid_len (B x 1) - int - valiud length of each sequence
+                    logits (B x T x NUM_LABELS) - logits for each token and each tag
+        """
+        for subwords, gold_tags, tokens, valid_len in dataloader:
+            self.model.train(is_train)
+
+            if torch.cuda.is_available():
+                subwords = subwords.cuda()
+                gold_tags = gold_tags.cuda()
+
+            if is_train:
+                self.optimizer.zero_grad()
+                logits = self.model(subwords)
+            else:
+                with torch.no_grad():
+                    logits = self.model(subwords)
+
+            yield subwords, gold_tags, tokens, valid_len, logits
+
+    def segments_to_file(self, segments, filename):
+        """
+        Write segments to file
+        :param segments: [List[Nested.data.dataset.Token]] - list of list of tokens
+        :param filename: str - output filename
+        :return: None
+        """
+        with open(filename, "w") as fh:
+            results = "\n\n".join(["\n".join([t.__str__() for t in segment]) for segment in segments])
+            fh.write("Token\tGold Tag\tPredicted Tag\n")
+            fh.write(results)
+            logging.info("Predictions written to %s", filename)
+
+    def save(self):
+        """
+        Save model checkpoint
+        :return:
+        """
+        filename = os.path.join(
+            self.output_path,
+            "checkpoints",
+            "checkpoint_{}.pt".format(self.current_epoch),
+        )
+
+        checkpoint = {
+            "model": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "epoch": self.current_epoch
+        }
+
+        logger.info("Saving checkpoint to %s", filename)
+        torch.save(checkpoint, filename)
+
+    def load(self, checkpoint_path):
+        """
+        Load model checkpoint
+        :param checkpoint_path: str - path/to/checkpoints
+        :return: None
+        """
+        # checkpoint_path = natsort.natsorted(glob.glob(f"{checkpoint_path}/checkpoint_*.pt"))
+        checkpoint_path = natsort.natsorted(checkpoint_path)
+        # checkpoint_path = checkpoint_path[-1]
+
+        logger.info("Loading checkpoint %s", checkpoint_path)
+
+        device = None if torch.cuda.is_available() else torch.device('cpu')
+        # checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        repo_path = snapshot_download(repo_id="SinaLab/Nested")
+
+        model_file = os.path.join(repo_path, "checkpoints", "checkpoint_2.pt")
+
+        checkpoint = torch.load(model_file, map_location=device, weights_only=False)
+
+        self.model.load_state_dict(checkpoint["model"])