Upload 2 files

model.py

@@ -0,0 +1,385 @@
+from transformers import AutoTokenizer, AutoModel
+from datasets import load_dataset, Dataset, concatenate_datasets
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import (
+    classification_report,
+    confusion_matrix,
+    accuracy_score,
+    precision_score,
+)
+from sklearn.ensemble import RandomForestClassifier
+from xgboost import XGBClassifier
+import torch.nn as nn
+import torchmetrics
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+import numpy as np
+import pandas as pd
+import os
+import pickle
+import argparse
+from torch_train import TorchTrain
+from utilities import get_simple_logger
+
+FILE_DIR = os.path.dirname(os.path.realpath(__file__))
+DATA_DIR = os.path.join(FILE_DIR, "data")
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+random_state = 42
+# set random state
+np.random.seed(random_state)
+torch.manual_seed(random_state)
+
+
+class PDFDataLoader:
+    """A class that can be used to load the data to torch model. This will be used in the `PDFDataSet` class to create the final datasets."""
+
+    def __init__(self, df):
+        self.df = df
+
+    def __getitem__(self, idx):
+        """Gets the `idx` embedding and labels, converts them to the required format and returns them."""
+        row = self.df[idx]
+        embeddings = row["embeddings"]
+        label = row["label"]
+        # convert to torch int
+        label = np.array(label)
+        # add extra dimension to label
+        label = np.expand_dims(label, axis=0)
+        embeddings = torch.from_numpy(np.array(embeddings)).float()
+        return embeddings.to(device), torch.from_numpy(label).to(device).float()
+
+    def __len__(self):
+        return len(self.df)
+
+
+class PDFDataSet:
+    def __init__(
+        self,
+        data_dir=DATA_DIR,
+        fraction_test_data_in_train=0.2,
+        model_ckpt="encoder",
+    ) -> None:
+        self.data_dir = data_dir
+        self.fraction_test_data_in_train = fraction_test_data_in_train
+        self.model_ckpt = model_ckpt
+        tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
+        encoding_model = AutoModel.from_pretrained(model_ckpt)
+        encoding_model = encoding_model.to(device)
+        encoding_model = encoding_model.eval()
+        self.encoding_model = encoding_model
+        self.tokenizer = tokenizer
+        self.logger = get_simple_logger("pdf_dataset")
+
+    def create_datasets(self):
+        train_data_path = os.path.join(FILE_DIR, self.data_dir, "train.csv")
+        test_data_path = os.path.join(FILE_DIR, self.data_dir, "test.csv")
+        df = pd.read_csv(train_data_path)
+        test_df = pd.read_csv(test_data_path)
+        train_df, validation_df = train_test_split(df, test_size=0.3, random_state=42)
+        if self.fraction_test_data_in_train:
+            self.logger.info(
+                f"Adding {self.fraction_test_data_in_train} fraction of test dataset to the training set."
+            )
+            test_df, test_df_for_training = train_test_split(
+                test_df, test_size=self.fraction_test_data_in_train, random_state=42
+            )
+            train_df = pd.concat([train_df, test_df_for_training])
+
+        train_dataset = Dataset.from_pandas(train_df)
+        validation_dataset = Dataset.from_pandas(validation_df)
+        test_dataset = Dataset.from_pandas(test_df)
+        return train_dataset, validation_dataset, test_dataset
+
+    def mean_pooling(self, model_output, attention_mask):
+        token_embeddings = model_output[
+            0
+        ]  # First element of model_output contains all token embeddings
+        input_mask_expanded = (
+            attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        )
+        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
+            input_mask_expanded.sum(1), min=1e-9
+        )
+
+    def sentences_to_embedding(self, sentences):
+        # Tokenize sentences
+        encoded_input = self.tokenizer(
+            sentences, padding=True, truncation=True, return_tensors="pt"
+        )
+        sentence_embeddings = self.mean_pooling(
+            self.encoding_model(**encoded_input), encoded_input["attention_mask"]
+        )
+        sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
+        # remove last dimension
+        sentence_embeddings = sentence_embeddings.squeeze()
+        return sentence_embeddings.detach()
+
+    def get_embeddings(self, row):
+        return {
+            "embeddings": self.sentences_to_embedding(
+                sentences=row["content"],
+            )
+        }
+
+    def create_embeddings(self):
+        train_dataset, validation_dataset, test_dataset = self.create_datasets()
+        train_dataset = train_dataset.map(self.get_embeddings)
+        validation_dataset = validation_dataset.map(self.get_embeddings)
+        test_dataset = test_dataset.map(self.get_embeddings)
+        return train_dataset, validation_dataset, test_dataset
+
+
+class PDFModel(nn.Module):
+    def __init__(self, input_size, hidden_sizes, output_size):
+        super(PDFModel, self).__init__()
+        self.seq_model = nn.Sequential()
+        for i, hidden_size in enumerate(hidden_sizes):
+            self.seq_model.add_module(f"linear_{i}", nn.Linear(input_size, hidden_size))
+            self.seq_model.add_module(f"relu_{i}", nn.ReLU())
+            input_size = hidden_size
+        self.last_layer = nn.Linear(input_size, output_size)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        seq_out = self.seq_model(x)
+        out = self.last_layer(seq_out)
+        return self.sigmoid(out)
+
+
+def evaluate_model(y_true, y_pred, model_name, split="train"):
+    accuracy = accuracy_score(y_true, y_pred)
+    precision = precision_score(y_true, y_pred)
+    classification_report_ = classification_report(y_true, y_pred)
+    print("------" * 10)
+    print(f"Evaluating for the model: {model_name} for {split} dataset...")
+    print(f"Accuracy: {accuracy}")
+    print(f"Precision: {precision}")
+    print(classification_report_)
+    print("------" * 10)
+
+
+def train_dl_model(
+    train_data,
+    validation_data,
+    epochs=30,
+    input_shape=384,
+    hidden_sizes=[32, 16],
+):
+    model = PDFModel(input_size=input_shape, hidden_sizes=hidden_sizes, output_size=1)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+    loss_fn = nn.BCELoss()
+    accuracy = torchmetrics.Accuracy(
+        task="binary", num_classes=2, threshold=0.5, average="macro"
+    )
+    precision = torchmetrics.Precision(task="binary", average="macro")
+    metrics = {
+        "accuracy": accuracy,
+        "precision": precision,
+    }
+    scheduler = CosineAnnealingLR(optimizer, T_max=10, eta_min=0.0001)
+    tt = TorchTrain(model, optimizer, loss_fn, metrics=metrics, scheduler=scheduler)
+    history = tt.fit(train_data, validation_data, verbose=True, epochs=epochs)
+    return history, model
+
+
+def evaluate_models(fraction_test_data_in_train=0.1):
+    print("Creating Embeddings...")
+    ds = PDFDataSet(fraction_test_data_in_train=fraction_test_data_in_train)
+    train_dataset, validation_dataset, test_dataset = ds.create_embeddings()
+    print("Done\n")
+
+    print("Training DL Model")
+    # Create dataset for DL models:
+    BATCH_SIZE = 8
+    train_dataloader = PDFDataLoader(train_dataset)
+    validation_dataloader = PDFDataLoader(validation_dataset)
+    test_dataloader = PDFDataLoader(test_dataset)
+
+    train_data = DataLoader(train_dataloader, batch_size=BATCH_SIZE, shuffle=True)
+    validation_data = DataLoader(
+        validation_dataloader,
+        batch_size=BATCH_SIZE,
+        shuffle=True,
+    )
+    test_data = DataLoader(test_dataloader, batch_size=BATCH_SIZE, shuffle=True)
+    for X, y in train_data:
+        input_shape = int(X.shape[1])
+        output_shape = int(y.shape[1])
+        break
+    epochs = 30
+    hidden_sizes = [32, 16]
+    history, model = train_dl_model(
+        train_data=train_data,
+        validation_data=validation_data,
+        epochs=epochs,
+        hidden_sizes=hidden_sizes,
+    )
+    print("Done\n")
+    print("Evaluating DL Model")
+    y_test_pred = model(torch.from_numpy(np.array(test_dataset["embeddings"])).float())
+    y_test_pred = y_test_pred.detach().numpy()
+    y_test_pred = np.where(y_test_pred > 0.5, 1, 0)
+    evaluate_model(
+        y_true=test_dataset["label"],
+        y_pred=y_test_pred,
+        model_name="DL Model",
+        split="test",
+    )
+    print("Done\n")
+
+    # ML Models
+    print("Training and evaluating ML Models.")
+    X_train = train_dataset["embeddings"]
+    y_train = train_dataset["label"]
+    X_validation = validation_dataset["embeddings"]
+    y_validation = validation_dataset["label"]
+    X_test = test_dataset["embeddings"]
+    y_test = test_dataset["label"]
+    rfc_best_params = {
+        "max_depth": 23,
+        "max_features": "log2",
+        "n_estimators": 469,
+    }
+
+    xgb_best_params = {
+        "max_depth": 25,
+        "n_estimators": 372,
+        "learning_rate": 0.2522824287799319,
+    }
+    print("Fitting RandomForest")
+    rfc = RandomForestClassifier(**rfc_best_params)
+    rfc.fit(X_train, y_train)
+    evaluate_model(
+        y_true=y_train,
+        y_pred=rfc.predict(X_train),
+        model_name="RandomForest",
+        split="train",
+    )
+    evaluate_model(
+        y_true=y_validation,
+        y_pred=rfc.predict(X_validation),
+        model_name="RandomForest",
+        split="validation",
+    )
+    evaluate_model(
+        y_true=y_test,
+        y_pred=rfc.predict(X_test),
+        model_name="RandomForest",
+        split="test",
+    )
+
+    print("Fitting XGBoost")
+    xgb = XGBClassifier(**xgb_best_params)
+    xgb.fit(X_train, y_train)
+    evaluate_model(
+        y_true=y_train,
+        y_pred=xgb.predict(X_train),
+        model_name="XGBoost",
+        split="train",
+    )
+    evaluate_model(
+        y_true=y_validation,
+        y_pred=xgb.predict(X_validation),
+        model_name="XGBoost",
+        split="validation",
+    )
+    evaluate_model(
+        y_true=y_test,
+        y_pred=xgb.predict(X_test),
+        model_name="XGBoost",
+        split="test",
+    )
+    print("All Done")
+
+
+def train_and_save_final_model(model_save_path="final_model.pkl"):
+    """This method creats and save the final model. The final model has the following characterstics:
+
+    - It is a RandomForestClassifier trained on all the training data and 10% of the test data. 10% of the test data. The 10% of test data is necessary as the distribution of the test data is very different from the training data.
+    - Since 10% of test data is used while training, this data is not used while claculating the final accuracy of the model, which is 100%.
+
+    Parameters
+    ----------
+    model_save_path : str, optional
+        The path to save the final model, by default "final_model.pkl"
+    Returns
+    -------
+    None
+    Examples
+    --------
+    >>> train_and_save_final_model()
+    >>> train_and_save_final_model(model_save_path="final_model.pkl")
+    """
+    print("Creating Embeddings...")
+    model_save_path = os.path.join(FILE_DIR, model_save_path)
+    ds = PDFDataSet(fraction_test_data_in_train=0.1)
+    train_dataset, validation_dataset, test_dataset = ds.create_embeddings()
+    train_dataset = concatenate_datasets([train_dataset, validation_dataset])
+    X_train = train_dataset["embeddings"]
+    X_test = test_dataset["embeddings"]
+    y_train = train_dataset["label"]
+    y_test = test_dataset["label"]
+
+    print("Training and evaluating the model...")
+    rfc_best_params = {
+        "max_depth": 23,
+        "max_features": "log2",
+        "n_estimators": 469,
+    }
+    rfc_model = RandomForestClassifier(**rfc_best_params)
+    rfc_model.fit(X_train, y_train)
+    evaluate_model(
+        y_true=y_train,
+        y_pred=rfc_model.predict(X_train),
+        model_name="Final Model",
+        split="train",
+    )
+    evaluate_model(
+        y_true=y_test,
+        y_pred=rfc_model.predict(X_test),
+        model_name="Final Model",
+        split="test",
+    )
+
+    print("Saving the model...")
+    with open(model_save_path, "wb") as f:
+        pickle.dump(rfc_model, f)
+    print(f"Model saved to: {model_save_path}")
+
+
+def main(args):
+    task = args.task
+    if task == "train":
+        model_save_path = args.model_save_path
+        train_and_save_final_model(model_save_path=model_save_path)
+    elif task == "evaluate":
+        fraction_test_data_in_train = args.fraction
+        evaluate_models(fraction_test_data_in_train)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train and evaluate models")
+    parser.add_argument(
+        "--task",
+        type=str,
+        choices=["train", "evaluate"],
+        required=True,
+        help="Whether to train and save the best model or evaluate all the models.",
+    )
+    parser.add_argument(
+        "--fraction",
+        type=float,
+        default=0.1,
+        help="Fraction of test data in train dataset",
+    )
+    parser.add_argument(
+        "--model_save_path",
+        type=str,
+        default="final_model.pkl",
+        help="Path to save the final model",
+    )
+    args = parser.parse_args()
+    main(args)

torch_train.py

@@ -0,0 +1,543 @@
+import torch
+import numpy as np
+
+
+class TorchTrain:
+    """A class for training a model in PyTorch.
+
+    Parameters
+    -----------
+        model (torch.nn.Module): The PyTorch model to train.
+        optimizer (torch.optim.Optimizer): The optimizer to use for training.
+        loss_function (callable): The loss function to use for training.
+        metrics (dict or callable, optional): The metrics to evaluate during training.
+            If a dictionary, the keys are the metric names and the values are functions that
+            take in `yhat` and `y` and return a metric value. If a callable, it should take
+            in `yhat` and `y` and return a metric value. Defaults to None.
+
+    Attributes
+    -----------
+        DEVICE (torch.device): The device to use for training (cuda if available, cpu otherwise).
+        model (torch.nn.Module): The PyTorch model being trained.
+        optimizer (torch.optim.Optimizer): The optimizer being used for training.
+        loss_function (callable): The loss function being used for training.
+        metrics (dict or callable): The metrics being evaluated during training.
+        metrics_evaluated (dict): The metrics evaluated during training.
+        train_loss (float): The average training loss.
+        test_loss (float): The average test loss.
+        train_iteration (int): The number of training iterations.
+        test_iteration (int): The number of test iterations.
+        train_metrics (dict): The metrics evaluated on the training data.
+        test_metrics (dict): The metrics evaluated on the test data.
+    """
+
+    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def __init__(
+        self,
+        model,
+        optimizer,
+        loss_function,
+        metrics=None,
+        scheduler=None,
+        task_type="classification",
+    ) -> None:
+        """Initialize the TorchTrain object.
+
+        Parameters
+        -----------
+        model : torch.nn.Module
+            The PyTorch model to train.
+        optimizer : torch.optim.Optimizer
+            The optimizer to use for training.
+        loss_function : callable
+            The loss function to use for training.
+        metrics : dict or callable, optional
+            The metrics to evaluate during training. If a dictionary, the keys are the metric names
+            and the values are functions that take in `yhat` and `y` and return a metric value.
+            If a callable, it should take in `yhat` and `y` and return a metric value. Defaults to None.
+        scheduler : torch.optim.lr_scheduler, optional
+            The learning rate scheduler to use for training. Defaults to None.
+        """
+        self.model = model
+        self.model.to(self.DEVICE)
+        self.optimizer = optimizer
+        self.loss_function = loss_function
+        self.metrics = self.__preprocess_metrics(metrics)
+        self.scheduler = scheduler
+        self.metrics_evaluated = {}
+        self.train_loss = 0
+        self.test_loss = 0
+        self.train_iteration = 0
+        self.test_iteration = 0
+        self.train_metrics = {}
+        self.test_metrics = {}
+        self.history = {}
+        self.train_loss_all = []
+        self.test_loss_all = []
+        self.train_metrics_all = []
+        self.test_metrics_all = []
+        self.__train_scaled = False
+        self.__test_scaled = False
+        self.task_type = task_type
+
+    def __preprocess_metrics(self, metrics):
+        """Preprocesses the given metrics"""
+        if metrics is None:
+            return {}
+        if isinstance(metrics, dict):
+            return {key.title(): value for key, value in metrics.items()}
+        else:
+            raise TypeError(
+                "Metrics should be a dictionary of metrics or a function which takes yhat, y"
+            )
+
+    def __scale_matrices(self, loss, metrics, type="train"):
+        """Scales the loss and metrics
+
+        Parameters
+        -----------
+        loss : float
+            The loss to scale
+        metrics : dict
+            The metrics to scale
+        type : str, optional
+            The type of scaling to do, either "train" or "test", by default "train"
+
+        Returns
+        --------
+        loss : float
+            The scaled loss
+        metrics : dict
+            The scaled metrics
+        """
+        if type == "train" and not self.__train_scaled:
+            scale = self.train_iteration
+            self.__train_scaled = True
+        elif type == "test" and not self.__test_scaled:
+            scale = self.test_iteration
+            self.__test_scaled = True
+        else:
+            return loss, metrics
+        loss /= scale
+        for key in metrics:
+            metrics[key] /= scale
+        return loss, metrics
+
+    def __reset_counters(self):
+        """Resets all the counters and loss objects for a new epoch"""
+        self.train_loss, self.train_metrics = self.__scale_matrices(
+            self.train_loss, self.train_metrics, type="train"
+        )
+
+        self.test_loss, self.test_metrics = self.__scale_matrices(
+            self.test_loss, self.test_metrics, type="test"
+        )
+
+        self.train_loss_all.append(self.train_loss)
+        self.train_loss = 0
+
+        self.test_loss_all.append(self.test_loss)
+        self.test_loss = 0
+
+        self.train_iteration = 0
+        self.test_iteration = 0
+
+        self.train_metrics_all.append(self.train_metrics)
+        self.train_metrics = {}
+
+        self.test_metrics_all.append(self.test_metrics)
+        self.test_metrics = {}
+        self.__train_scaled = False
+        self.__test_scaled = False
+
+    @property
+    def loss(self):
+        """Returns the training loss"""
+        return self.train_loss_all[-1]
+
+    def __create_history(self):
+        """Creates the history dictionary"""
+        history = {
+            "train_loss": self.train_loss_all,
+            "val_loss": self.test_loss_all,
+        }
+        for key, value in self.metrics.items():
+            history[f"train_{key.lower()}"] = []
+            history[f"val_{key.lower()}"] = []
+
+        for item in self.train_metrics_all:
+            for key, value in item.items():
+                history[f"train_{key.lower()}"].append(value)
+
+        for item in self.test_metrics_all:
+            for key, value in item.items():
+                history[f"val_{key.lower()}"].append(value)
+        return history
+
+    def __parse_val(self, val):
+        """Parses the given value to a float"""
+        if isinstance(val, torch.Tensor):
+            val = val.item()
+        elif isinstance(val, np.ndarray):
+            val = float(val)
+        elif isinstance(val, (int, float)):
+            pass
+        else:
+            raise TypeError(
+                f"The given Metric function should return a tensor, numpy array, int, or float.\n\
+                    Got {type(val)}"
+            )
+        return val
+
+    def _train_step(self, x, y):
+        """Perform a single training step.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The input tensor.
+        y : torch.Tensor
+            The target tensor.
+
+        Returns
+        -------
+        tuple
+            A tuple containing the loss and the predicted output tensor.
+        """
+        self.model.train()
+        yhat = self.model(x)
+        l = self.loss_function(yhat, y)
+        self.optimizer.zero_grad()
+        l.backward()
+        self.optimizer.step()
+        self.train_iteration += 1
+        return l.item(), yhat
+
+    def _test_step(self, x, y):
+        """Perform a single testing step.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The input tensor.
+        y : torch.Tensor
+            The target tensor.
+
+        Returns
+        -------
+        tuple
+            A tuple containing the loss and the predicted output tensor.
+        """
+        self.model.eval()
+        with torch.inference_mode():
+            yhat = self.model(x)
+            l = self.loss_function(yhat, y)
+            self.test_iteration += 1
+        return l.item(), yhat
+
+    def predict(self, x):
+        """Make predictions on a batch of data.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            The predicted output tensor.
+        """
+        self.model.eval()
+        yhat = self.model(x)
+        if self.task_type == "classification":
+            if len(yhat.shape) == 1:
+                # round
+                yhat = torch.round(yhat)
+                yhat = yhat.unsqueeze(1)
+            else:
+                yhat = torch.argmax(yhat, dim=1)
+
+        return yhat
+
+    def __calculate_metrics(self, yhat, y):
+        """Calculate the metrics for a batch of data.
+
+        Parameters
+        ----------
+        yhat : torch.Tensor
+            The predicted output tensor.
+        y : torch.Tensor
+            The target tensor.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the values of the metrics.
+        """
+        metrics = {}
+        for key, metric in self.metrics.items():
+            val = metric(yhat, y)
+            if isinstance(val, torch.Tensor):
+                val = val.item()
+            elif isinstance(val, np.ndarray):
+                val = float(val)
+            elif isinstance(val, (int, float)):
+                pass
+            else:
+                raise TypeError(
+                    f"Metric {key} should return a tensor, numpy array, int, or float"
+                )
+            metrics[key] = val
+        self.metrics_evaluated = metrics
+        return metrics
+
+    def __progress_bar(self, cur_iter, all_iter):
+        """Creates a progress bar showing the progress of the current batch.
+
+        Parameters
+        ----------
+        cur_iter : int
+            The current batch number.
+        all_iter : int
+            The total number of batches.
+
+        Returns
+        -------
+        str
+            The progress bar, in the form of "10/100[====----]".
+        """
+        len_progress_bar = 20
+        progress = int((cur_iter + 1) / all_iter * len_progress_bar)
+        progress_bar = "=" * progress + "-" * (len_progress_bar - progress)
+        return f"[{progress_bar}]"
+
+    def progress(self, cur_iter, all_iter, loss, metrics, on="train"):
+        """Prints a progress bar showing the progress of the current batch.
+
+        Parameters
+        ----------
+        cur_iter : int
+            The current batch number.
+        all_iter : int
+            The total number of batches.
+        loss : float
+            The current loss. Should be averaged over all batches.
+        metrics : dict
+            The metrics evaluated on the current batch.
+        on : str, optional
+            Whether the progress bar is for the training or testing data. Defaults to "train".
+
+        Returns
+        -------
+        str
+            The progress bar, in the form of "10/100[====----]".
+
+        Notes
+        -----
+        The progress bar shows the progress of the current batch as a bar of equal signs ("=") and
+        hyphens ("-"). The length of the bar is fixed at 20 characters. The current batch number
+        and total number of batches are displayed at the beginning of the progress bar. The current
+        loss and any metrics evaluated on the current batch are displayed at the end of the progress
+        bar.
+        """
+        # len_progress_bar = 20
+        # progress = int((cur_iter + 1) / all_iter * len_progress_bar)
+        # progress_bar = "=" * progress + "-" * (len_progress_bar - progress)
+        progress_bar = self.__progress_bar(cur_iter=cur_iter, all_iter=all_iter)
+
+        if on.lower() == "train":
+            iteration = self.train_iteration
+            prefix = f"Epoch {(self.current_epoch+1):2d}/{self.epochs:2d} Batch "
+        else:
+            iteration = self.test_iteration
+            prefix = "Epoch "
+
+        text = f"{prefix}{cur_iter:>4d}/{all_iter:>4d}{progress_bar} {on.title()} loss: {loss/iteration:.4f}"
+        for metric_name, metric_value in metrics.items():
+            text += f" | {on.title()} {metric_name}: {metric_value/iteration:.4f}"
+
+        return text
+
+    def update_metrics(self, cur_metrics, new_metrics):
+        """Update the metrics with the values for a new batch of data.
+
+        Parameters
+        ----------
+        cur_metrics : dict
+            The current values of the metrics.
+        new_metrics : dict
+            The values of the metrics for a new batch of data.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the updated values of the metrics.
+        """
+        for key, value in new_metrics.items():
+            if key not in cur_metrics:
+                cur_metrics[key] = value
+            else:
+                cur_metrics[key] += value
+        return cur_metrics
+
+    def fit(
+        self,
+        train_loader,
+        validation_data_loader=None,
+        epochs=1,
+        verbose=True,
+        train_steps_per_epoch=None,
+        validation_steps_per_epoch=None,
+    ):
+        """Fit the PyTorch model.
+
+        Parameters
+        ----------
+        train_loader : torch.utils.data.DataLoader
+            The data loader for the training data.
+        validation_data_loader : torch.utils.data.DataLoader, optional
+            The data loader for the test data. Defaults to None.
+        epochs : int, optional
+            The number of epochs to train for. Defaults to 1.
+        verbose : bool, optional
+            Whether to print the training progress during training. Defaults to True.
+        train_steps_per_epoch : int, optional
+            The number of batches to train on per epoch. Defaults to None.
+        validation_steps_per_epoch : int, optional
+            The number of batches to test on per epoch. Defaults to None.
+
+        Returns
+        -------
+        None
+
+        Examples
+        --------
+        >>> model = MyModel()
+        >>> optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+        >>> loss_function = nn.CrossEntropyLoss()
+        >>> scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)
+        >>> train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
+        >>> validation_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)
+        >>> trainer = TorchTrain(model, optimizer, loss_function, scheduler=scheduler)
+        >>> trainer.fit(train_loader, validation_data_loader=validation_data_loader, epochs=10, verbose=True)
+        """
+        self.epochs = epochs
+        if train_steps_per_epoch is None:
+            train_steps_per_epoch = len(train_loader)
+        if validation_data_loader is not None:
+            if validation_steps_per_epoch is None:
+                validation_steps_per_epoch = len(validation_data_loader)
+
+        for epoch in range(epochs):
+            self.current_epoch = epoch
+            for i, (x, y) in enumerate(train_loader):
+                x = x.to(self.DEVICE)
+                if isinstance(y, list) or isinstance(y, tuple):
+                    y = [y_.to(self.DEVICE) for y_ in y]
+                else:
+                    y = y.to(self.DEVICE)
+
+                train_loss, yhat = self._train_step(x, y)
+                self.train_loss += train_loss
+                metrics = self.__calculate_metrics(yhat, y)
+                self.train_metrics = self.update_metrics(self.train_metrics, metrics)
+
+                b_progress = self.progress(
+                    i + 1,
+                    train_steps_per_epoch,
+                    self.train_loss,
+                    self.train_metrics,
+                    on="train",
+                )
+                if i == train_steps_per_epoch - 1:
+                    print(b_progress)
+                    break
+                else:
+                    if verbose:
+                        print(b_progress, end="\r")
+            if validation_data_loader is not None:
+                for i, (x, y) in enumerate(validation_data_loader):
+                    x = x.to(self.DEVICE)
+                    if isinstance(y, list) or isinstance(y, tuple):
+                        y = [y_.to(self.DEVICE) for y_ in y]
+                    else:
+                        y = y.to(self.DEVICE)
+                    test_loss, yhat = self._test_step(x, y)
+                    self.test_loss += test_loss
+                    metrics = self.__calculate_metrics(yhat, y)
+                    self.test_metrics = self.update_metrics(self.test_metrics, metrics)
+                    if i == validation_steps_per_epoch - 1:
+                        break
+                test_progress = self.progress(
+                    epoch + 1,
+                    epochs,
+                    self.test_loss,
+                    self.test_metrics,
+                    on="test",
+                )
+                print(test_progress)
+            self.__reset_counters()
+            if self.scheduler is not None:
+                self.scheduler.step()
+            if verbose and self.scheduler is not None:
+                print(f"New Learning rate: {self.scheduler.get_last_lr()[0]:.6f}")
+
+        return self.__create_history()
+
+    def save(self, path):
+        """Save the model to a file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the file to save the model to.
+        """
+        torch.save(self.model.state_dict(), path)
+
+    def load(self, path):
+        """Load the model from a file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the file to load the model from.
+        """
+        self.model.load_state_dict(torch.load(path))
+
+    def evaluate(self, data_loader, metric):
+        """Evaluate the model on a data loader and the given metric.
+
+        Parameters
+        ----------
+        data_loader : torch.utils.data.DataLoader
+            The data loader to evaluate the model on.
+        metric : function
+            The metric to evaluate the model with.
+
+        Returns
+        -------
+        float
+            The score of the model on the given metric.
+        """
+        running_score = 0
+        data_length = len(data_loader)
+        for i, (x, y) in enumerate(data_loader):
+            progress_bar = self.__progress_bar(i, data_length)
+            x = x.to(self.DEVICE)
+            if isinstance(y, list) or isinstance(y, tuple):
+                y = [y_.to(self.DEVICE) for y_ in y]
+            else:
+                y = y.to(self.DEVICE)
+
+            yhat = self.model(x)
+            yhat = torch.round(yhat)
+            score = metric(y, yhat)
+            score = self.__parse_val(score)
+            running_score += score
+
+            progress_bar = f"{i+1}/{data_length}" + progress_bar
+            progress_bar += f" Score: {(running_score/(i+1)):4f}"
+            print(progress_bar, end="\r")
+        return running_score / (len(data_loader))