hotfix

Artur-Galstyan · Apr 8, 2024 · 0575a8f · 0575a8f
1 parent 2059dca
commit 0575a8f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/jaxonloader/dataset.py b/jaxonloader/dataset.py
@@ -30,6 +30,12 @@ def __len__(self) -> int:
     def __getitem__(self, idx):
         return self.data[idx]
 
+    def split(self, ratio: float) -> tuple["SingleArrayDataset", "SingleArrayDataset"]:
+        split = int(len(self.data) * ratio)
+        return SingleArrayDataset(self.data[:split]), SingleArrayDataset(
+            self.data[split:]
+        )
+
 
 class DataTargetDataset(JaxonDataset):
     def __init__(self, data: NDArray, target: NDArray):
@@ -43,3 +49,9 @@ def __len__(self) -> int:
 
     def __getitem__(self, idx):
         return self.data[idx], self.target[idx]
+
+    def split(self, ratio: float) -> tuple["DataTargetDataset", "DataTargetDataset"]:
+        split = int(len(self.data) * ratio)
+        return DataTargetDataset(
+            self.data[:split], self.target[:split]
+        ), DataTargetDataset(self.data[split:], self.target[split:])
diff --git a/noxfile.py b/noxfile.py
@@ -3,6 +3,7 @@
 
 @nox.session
 def tests(session):
+    session.install("uv")
     session.install("pip")
     session.run("uv", "pip", "install", ".[dev]")
     session.run("uv", "pip", "install", "-e", ".")