Merge branch 'main' into add_options_mixed

Author: MMathisLab

.github/workflows/build.yml

@@ -23,10 +23,9 @@ jobs:
           - os: ubuntu-latest
             python-version: 3.8
             torch-version: 1.9.0
-          # TODO(stes): Include at a later stage
-          #- os: windows-latest
-          #  torch-version: 2.0.0
-          #  python-version: "3.10"
+          - os: windows-latest
+            torch-version: 2.0.0
+            python-version: "3.10"
           #- os: macos-latest
           #  torch-version: 2.0.0
           #  python-version: "3.10"

cebra/integrations/sklearn/cebra.py

@@ -1161,14 +1161,18 @@ def fit(
 
             >>> import cebra
             >>> import numpy as np
+            >>> import tempfile
+            >>> from pathlib import Path
+            >>> tmp_file = Path(tempfile.gettempdir(), 'cebra.pt')
             >>> dataset =  np.random.uniform(0, 1, (1000, 20))
             >>> dataset2 =  np.random.uniform(0, 1, (1000, 40))
             >>> cebra_model = cebra.CEBRA(max_iterations=10)
             >>> cebra_model.fit(dataset)
             CEBRA(max_iterations=10)
-            >>> cebra_model.save('/tmp/foo.pt')
+            >>> cebra_model.save(tmp_file)
             >>> cebra_model.fit(dataset2, adapt=True)
             CEBRA(max_iterations=10)
+            >>> tmp_file.unlink()
         """
         if adapt and sklearn_utils.check_fitted(self):
             self._adapt_fit(X,
@@ -1332,11 +1336,15 @@ def save(self,
 
             >>> import cebra
             >>> import numpy as np
+            >>> import tempfile
+            >>> from pathlib import Path
+            >>> tmp_file = Path(tempfile.gettempdir(), 'test.jl')
             >>> dataset =  np.random.uniform(0, 1, (1000, 30))
             >>> cebra_model = cebra.CEBRA(max_iterations=10)
             >>> cebra_model.fit(dataset)
             CEBRA(max_iterations=10)
-            >>> cebra_model.save('/tmp/foo.pt')
+            >>> cebra_model.save(tmp_file)
+            >>> tmp_file.unlink()
 
         """
         if sklearn_utils.check_fitted(self):
@@ -1394,10 +1402,18 @@ def load(cls,
         Example:
 
             >>> import cebra
-            >>> import numpy as np
+            >>> import numpy as np        
+            >>> import tempfile
+            >>> from pathlib import Path
+            >>> tmp_file = Path(tempfile.gettempdir(), 'cebra.pt')
             >>> dataset =  np.random.uniform(0, 1, (1000, 20))
-            >>> loaded_model = cebra.CEBRA.load('/tmp/foo.pt')
+            >>> cebra_model = cebra.CEBRA(max_iterations=10)
+            >>> cebra_model.fit(dataset)
+            CEBRA(max_iterations=10)
+            >>> cebra_model.save(tmp_file)
+            >>> loaded_model = cebra.CEBRA.load(tmp_file)
             >>> embedding = loaded_model.transform(dataset)
+            >>> tmp_file.unlink()
 
         """
 

cebra/io.py

@@ -214,9 +214,11 @@ class FileKeyValueDataset:
 
         >>> import cebra.io
         >>> import joblib
-        >>> joblib.dump({'foo' : 42}, '/tmp/test.jl')
-        ['/tmp/test.jl']
-        >>> data = cebra.io.FileKeyValueDataset('/tmp/test.jl')
+        >>> import tempfile
+        >>> from pathlib import Path
+        >>> tmp_file = Path(tempfile.gettempdir(), 'test.jl')
+        >>> _ = joblib.dump({'foo' : 42}, tmp_file)
+        >>> data = cebra.io.FileKeyValueDataset(tmp_file) 
         >>> data.foo
         42
 
@@ -242,14 +244,14 @@ def __repr__(self):
         return f"{type(self).__name__}(keys=(\n  {sizes}\n))"
 
     def _iterate_items(self):
-        extension = self.path.split(".")[-1]
-        if extension in ["jl", "joblib"]:
+        extension = self.path.suffix
+        if extension in [".jl", ".joblib"]:
             dataset = joblib.load(self.path)
-        elif extension in ["h5", "hdf", "hdf5"]:
+        elif extension in [".h5", ".hdf", ".hdf5"]:
             raise NotImplementedError()
-        elif extension in ["pth", "pt"]:
+        elif extension in [".pth", ".pt"]:
             dataset = torch.load(self.path)
-        elif extension in ["npz"]:
+        elif extension in [".npz"]:
             dataset = np.load(self.path, allow_pickle=True)
         else:
             raise ValueError(f"Invalid file format: {extension} in {self.path}")

docs/source/usage.rst

@@ -494,14 +494,20 @@ The model will be saved as a ``.pt`` file.
 
 .. testcode::
 
+    import tempfile
+    from pathlib import Path
+
+    # create temporary file to save the model
+    tmp_file = Path(tempfile.gettempdir(), 'cebra.pt')
+
     cebra_model = cebra.CEBRA(max_iterations=10)
     cebra_model.fit(neural_data)
 
     # Save the model
-    cebra_model.save('/tmp/foo.pt')
+    cebra_model.save(tmp_file)
 
     # New session: load and use the model
-    loaded_cebra_model = cebra.CEBRA.load('/tmp/foo.pt')
+    loaded_cebra_model = cebra.CEBRA.load(tmp_file)
     embedding = loaded_cebra_model.transform(neural_data)
 
 
@@ -1221,10 +1227,11 @@ Putting all previous snippet examples together, we obtain the following pipeline
      cebra_model.fit(train_data, train_discrete_label, train_continuous_label)
 
      # 5. Save the model
-     cebra_model.save('/tmp/foo.pt')
+     tmp_file = Path(tempfile.gettempdir(), 'cebra.pt')
+     cebra_model.save(tmp_file)
 
      # 6. Load the model and compute an embedding
-     cebra_model = cebra.CEBRA.load('/tmp/foo.pt')
+     cebra_model = cebra.CEBRA.load(tmp_file)
      train_embedding = cebra_model.transform(train_data)
      valid_embedding = cebra_model.transform(valid_data)
      assert train_embedding.shape == (70, 8)