Switch to PyTorch Dataloader and HF datasets

Pull Request Title

Description

Wants to merge: vladb/py_dataloader into main

Type of change

• Bug fix
• New feature
• Enhancement
• Documentation update
• Other (specify right below)

Merge request commits

• Switch to PyTorch Dataloader

Related Issues

Screenshots or GIFs

Checklist

• I have tested the code with the changes manually.
• My code follows the project's style guidelines.
• I have documented my code for others to understand.
• I have updated documentation as needed (including README.md, code comments and doc strings).

Reviewer Guidelines

Additional Notes

@mentions

optimus/datasets/dataset_utils.py

@@ -64,6 +64,7 @@ def download_data(url: str, path: str) -> None:
         pass
 
 def tokenize_dataset(data, tokenizer):
+    """Tokenizes a list of strings"""
     start = time.time()
     print("Tokenizing dataset...")
 
@@ -74,4 +75,9 @@ def tokenize_dataset(data, tokenizer):
 
     print(f"Done. Took {time.time() - start:.2f}s.")
 
-    return data_tok
\ No newline at end of file
+    return data_tok
+
+def flatten_tokenized(data, seq_len):
+    data = torch.cat(data, dim=-1)
+
+    return data
\ No newline at end of file