Merge branch 'fix/full-dataset' 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

@@ -50,12 +50,12 @@ 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')
         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')
 
         bytes_array = np.frombuffer(self.bytes, dtype=np.uint8) # Zero-copy

src/utils/utils.py

@@ -1,3 +1,5 @@
+from os import path
+
 import torch
 from torch.utils.data import TensorDataset
 import matplotlib.pyplot as plt
@@ -14,7 +16,7 @@ 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 +25,12 @@ 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 load_data(path: str) -> bytes: