fix pkl file not loading in StaticDataLoader (#1896)

* fix pkl file not loading in StaticDataLoader

* resolve hard code

* resolve hard code

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
   "matplotlib",
   "jupyter",
   "nbconvert",
+  "pyarrow",
 ]
 
 [project.optional-dependencies]

qlib/data/dataset/loader.py

@@ -279,6 +279,9 @@ class StaticDataLoader(DataLoader, Serializable):
             )
             self._data.sort_index(inplace=True)
         elif isinstance(self._config, (str, Path)):
+            if str(self._config).strip().endswith(".parquet"):
+                self._data = pd.read_parquet(self._config, engine="pyarrow")
+            else:
                 with Path(self._config).open("rb") as f:
                     self._data = pickle.load(f)
         elif isinstance(self._config, pd.DataFrame):

qlib/rl/trainer/trainer.py

@@ -200,7 +200,7 @@ class Trainer:
 
         if ckpt_path is not None:
             _logger.info("Resuming states from %s", str(ckpt_path))
-            self.load_state_dict(torch.load(ckpt_path))
+            self.load_state_dict(torch.load(ckpt_path, weights_only=False))
         else:
             self.initialize()
 

tests/rl/test_trainer.py

@@ -194,7 +194,7 @@ def test_trainer_checkpoint():
     assert (output_dir / "002.pth").exists()
     assert os.readlink(output_dir / "latest.pth") == str(output_dir / "002.pth")
 
-    trainer.load_state_dict(torch.load(output_dir / "001.pth"))
+    trainer.load_state_dict(torch.load(output_dir / "001.pth", weights_only=False))
     assert trainer.current_iter == 1
     assert trainer.current_episode == 100