Merge branch 'portability' of github.ugent.be:ML/neural-compression into portability

src/args.py

@@ -19,7 +19,7 @@ def parse_arguments():
                              help="Which model to use")
     modelparser.add_argument("--model-load-path", type=str, required=False,
                              help="Filepath to the model to load")
-    modelparser.add_argument("--model-save-path", type=str, required=True,
+    modelparser.add_argument("--model-save-path", type=str, required=False,
                              help="Filepath to the model to save")
 
     fileparser = ArgumentParser(add_help=False)

src/dataset_loaders/Dataset.py

@@ -1,7 +1,9 @@
 from abc import abstractmethod, ABC
+from itertools import accumulate
 from os.path import join, curdir
 from typing import Callable
 
+import numpy as np
 import torch
 from torch import Tensor
 from torch.utils.data import Dataset as TorchDataset
@@ -49,25 +51,36 @@ class Dataset(TorchDataset, ABC):
         return len(self.dataset)
 
     def process_data(self):
+        self.chunk_offsets = self.get_offsets()
         if self.size == -1:
             # Just use the whole dataset
-            self.bytes = ''.join(tqdm(self.data, desc="Encoding data")).encode('utf-8', errors='replace')
+            self.bytes = ''.join(tqdm(self.data, desc="Encoding data", leave=False)).encode('utf-8', errors='replace')
         else:
             # Use only partition, calculate offsets
-            self.chunk_offsets = self.get_offsets()
-            self.bytes = ''.join(tqdm(self.data[:len(self.chunk_offsets)], desc="Encoding data")).encode('utf-8', errors='replace')
+            self.bytes = (''.join(tqdm(self.data[:len(self.chunk_offsets)], desc="Encoding data", leave=False))
+                          .encode('utf-8', errors='replace'))
 
-        self.tensor = torch.tensor(list(self.bytes), dtype=torch.long)
+        bytes_array = np.frombuffer(self.bytes, dtype=np.uint8)  # Zero-copy
+        self.tensor = torch.from_numpy(bytes_array).to(torch.long, non_blocking=True)
 
     def get_offsets(self):
         """
         Calculate for each chunk how many bytes came before it
         """
+        data = self.data
+        size = self.size
+
+        if size == -1:
+            return [0, *accumulate(tqdm(map(len, data), desc="Calculating offsets", leave=False, total=len(data)))]
+
         offsets = [0]
-        while len(offsets) <= len(self.data) and (self.size == -1 or offsets[-1] < self.size):
-            idx = len(offsets) - 1
-            offsets.append(offsets[idx] + len(self.data[idx]))
-        print(offsets)
+        total = 0
+        append = offsets.append
+        for chunk in tqdm(data):
+            if total >= size:
+                break
+            total += len(chunk)
+            append(total)
         return offsets
 
     def get_chunked_item(self, idx: int, offsets: list[int], context_length: int):

src/utils/utils.py

@@ -1,6 +1,8 @@
+from os import path
+
+import matplotlib.pyplot as plt
 import torch
 from torch.utils.data import TensorDataset
-import matplotlib.pyplot as plt
 
 
 def make_context_pairs(data: bytes, context_length: int) -> TensorDataset:
@@ -10,11 +12,13 @@ def make_context_pairs(data: bytes, context_length: int) -> TensorDataset:
     y = data[context_length:]
     return TensorDataset(x, y)
 
+
 def print_distribution(from_to: tuple[int, int], probabilities: list[float]):
     plt.hist(range(from_to[0], from_to[1]), weights=probabilities)
     plt.show()
 
-def print_losses(train_losses: list[float], validation_losses: list[float], show=False):
+
+def print_losses(train_losses: list[float], validation_losses: list[float], filename: str | None = None, show=False):
     plt.plot(train_losses, label="Training loss")
     plt.plot(validation_losses, label="Validation loss")
     plt.xlabel("Epoch")
@@ -23,7 +27,26 @@ def print_losses(train_losses: list[float], validation_losses: list[float], show
 
     if show:
         plt.show()
-    plt.savefig("losses.png")
+
+    if filename is None:
+        filename = path.join("results", "losses.png")
+
+    print(f"Saving losses to {filename}...")
+    plt.savefig(filename)
+
+
+def determine_device():
+    # NVIDIA GPUs (most HPC clusters)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    # Apple Silicon (macOS)
+    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return torch.device("mps")
+    # Intel GPUs (oneAPI)
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        return torch.device("xpu")
+    else:
+        return torch.device("cpu")
 
 
 def load_data(path: str) -> bytes: