0.14.40

lightdatamodule.py returns correct predictor

Author: bartzbeielstein

.gitignore

@@ -320,3 +320,4 @@ notebooks/00_spotPython_tests_files/libs/quarto-html/quarto.js
 notebooks/00_spotPython_tests_files/libs/quarto-html/tippy.css
 notebooks/00_spotPython_tests_files/libs/quarto-html/tippy.umd.min.js
 notebooks/00_spotPython_tests.html
+data.csv

pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spotpython"
-version = "0.14.38"
+version = "0.14.40"
 authors = [
   { name="T. Bartz-Beielstein", email="tbb@bartzundbartz.de" }
 ]

src/spotPython/data/lightdatamodule.py

@@ -298,5 +298,4 @@ def predict_dataloader(self) -> DataLoader:
         # print(f"LightDataModule: predict_dataloader(). batch_size: {self.batch_size}")
         # print(f"LightDataModule: predict_dataloader(). num_workers: {self.num_workers}")
         # apply fit_transform to the val data
-        return DataLoader(self.data_test, batch_size=self.batch_size, num_workers=self.num_workers)
         return DataLoader(self.data_predict, batch_size=len(self.data_predict), num_workers=self.num_workers)

test/test_lightdatamodule.py

@@ -1,4 +1,3 @@
-import pytest
 import torch
 from spotPython.data.lightdatamodule import LightDataModule
 from spotPython.data.csvdataset import CSVDataset
@@ -44,7 +43,3 @@ def test_light_data_module_sizes():
     assert len(data_module.data_train) == 3
     assert len(data_module.data_val) == 3
     assert len(data_module.data_test) == 6
-
-
-if __name__ == "__main__":
-    pytest.main(["-v", __file__])

test/test_lightdatamodule_class.py

@@ -0,0 +1,37 @@
+import pytest
+import torch
+from spotPython.data.lightdatamodule import LightDataModule
+from spotPython.data.csvdataset import CSVDataset
+
+
+class TestLightDataModule:
+    @pytest.fixture
+    def setup_data_module(self):
+        # Setup the dataset and data module as per the provided code snippet.
+        # Mock the data.csv file content
+        csv_content = """feature1,feature2,prognosis
+                        1,2,0
+                        3,4,1
+                        5,6,0
+                        7,8,1
+                        9,10,0
+                        11,12,1
+                        13,14,0
+                        15,16,1
+                        17,18,0
+                        19,20,1
+                        21,22,0
+                        23,24,1"""
+
+        with open("data.csv", "w") as f:
+            f.write(csv_content)
+
+        dataset = CSVDataset(csv_file="data.csv", target_column="prognosis", feature_type=torch.long)
+        data_module = LightDataModule(dataset=dataset, batch_size=5, test_size=0.5)
+        data_module.setup()
+        return data_module
+
+    def test_predict_set_size(self, setup_data_module):
+        data_module = setup_data_module
+        assert len(data_module.data_predict) == 6, "Expected predict set size to be 6,"
+        f"but got {len(data_module.data_predict)}"