fix transforms

cyclops/monitor/tester.py

@@ -717,33 +717,45 @@ def __init__(
                 batch_size=batch_size,
             )
             self.base_model.initialize()
-            self.transforms = partial(apply_transforms, transforms=transforms)
-            model_transforms = transforms
-            model_transforms.transforms = model_transforms.transforms + (
-                Lambdad(
-                    keys=("mask", "labels"),
-                    func=lambda x: np.array(x),
-                    allow_missing_keys=True,
-                ),
-            )
-            self.model_transforms = partial(
-                apply_transforms,
-                transforms=model_transforms,
+            self.base_model.save_model(
+                "saved_models/DetectronModule/pretrained_model.pt", log=False
             )
+            if transforms:
+                self.transforms = partial(apply_transforms, transforms=transforms)
+                model_transforms = transforms
+                model_transforms.transforms = model_transforms.transforms + (
+                    Lambdad(
+                        keys=("mask", "labels"),
+                        func=lambda x: np.array(x),
+                        allow_missing_keys=True,
+                    ),
+                )
+                self.model_transforms = partial(
+                    apply_transforms,
+                    transforms=model_transforms,
+                )
+            else:
+                self.transforms = None
+                self.model_transforms = None
         elif is_sklearn_model(base_model):
             self.base_model = wrap_model(base_model)
             self.base_model.save_model(
                 "saved_models/DetectronModule/pretrained_model.pkl", log=False
             )
             self.transforms = transforms
             self.model_transforms = transforms
-        elif isinstance(base_model, (PTModel, SKModel)):
+        elif isinstance(base_model, SKModel):
             self.base_model = base_model
             self.base_model.save_model(
                 "saved_models/DetectronModule/pretrained_model.pkl", log=False
             )
             self.transforms = transforms
             self.model_transforms = transforms
+        elif isinstance(base_model, PTModel):
+            self.base_model = base_model
+            self.base_model.save_model(
+                "saved_models/DetectronModule/pretrained_model.pt", log=False
+            )
         else:
             raise ValueError("base_model must be a PyTorch or sklearn model.")
 
@@ -774,9 +786,14 @@ def fit(self, X_s: Union[Dataset, DatasetDict, np.ndarray, TorchDataset]):
         for seed in range(self.num_runs):
             # train ensemble of for split 'p*'
             for e in range(1, self.ensemble_size + 1):
-                self.base_model.load_model(
-                    "saved_models/DetectronModule/pretrained_model.pkl", log=False
-                )
+                if is_pytorch_model(self.base_model.model):
+                    self.base_model.load_model(
+                        "saved_models/DetectronModule/pretrained_model.pt", log=False
+                    )
+                elif is_sklearn_model(self.base_model.model):
+                    self.base_model.load_model(
+                        "saved_models/DetectronModule/pretrained_model.pkl", log=False
+                    )
                 alpha = 1 / (len(X_s) * self.sample_size + 1)
                 if is_pytorch_model(self.base_model.model):
                     model = wrap_model(
@@ -896,9 +913,14 @@ def predict(self, X_t: Union[Dataset, DatasetDict, np.ndarray, TorchDataset]):
         for seed in range(self.num_runs):
             # train ensemble of for split 'p*'
             for e in range(1, self.ensemble_size + 1):
-                self.base_model.load_model(
-                    "saved_models/DetectronModule/pretrained_model.pkl", log=False
-                )
+                if is_pytorch_model(self.base_model.model):
+                    self.base_model.load_model(
+                        "saved_models/DetectronModule/pretrained_model.pt", log=False
+                    )
+                elif is_sklearn_model(self.base_model.model):
+                    self.base_model.load_model(
+                        "saved_models/DetectronModule/pretrained_model.pkl", log=False
+                    )
                 alpha = 1 / (len(X_t) * self.sample_size + 1)
                 if is_pytorch_model(self.base_model.model):
                     model = wrap_model(

docs/source/tutorials/diabetes_130/readmission_prediction_detectron.ipynb

@@ -514,7 +514,7 @@
     "    base_model=readmission_prediction_task.models[\"xgb_classifier\"],\n",
     "    feature_column=features_list,\n",
     "    transforms=preprocessor,\n",
-    "    splits_mapping={\"train\": \"train\", \"test\": \"test\"},\n",
+    "    splits_mapping={\"train\": \"train\", \"test\": \"validation\"},\n",
     "    sample_size=250,\n",
     "    num_runs=5,\n",
     "    ensemble_size=5,\n",
@@ -537,9 +537,9 @@
    "outputs": [],
    "source": [
     "results = tester.predict(\n",
-    "    X_t=DatasetDict({\"train\": dataset[\"train\"][\"train\"], \"test\": dataset[\"test\"]})\n",
+    "    X_t=DatasetDict({\"train\": dataset[\"train\"][\"train\"], \"validation\": dataset[\"test\"]})\n",
     ")\n",
-    "print(results[\"model_health\"])"
+    "print(results[\"data\"][\"model_health\"])"
    ]
   },
   {
@@ -575,9 +575,9 @@
     "model_health = []\n",
     "for data in test_data_list:\n",
     "    results = tester.predict(\n",
-    "        X_t=DatasetDict({\"train\": dataset[\"train\"][\"train\"], \"test\": data})\n",
+    "        X_t=DatasetDict({\"train\": dataset[\"train\"][\"train\"], \"validation\": data})\n",
     "    )\n",
-    "    model_health.append(results[\"model_health\"])"
+    "    model_health.append(results[\"data\"][\"model_health\"])"
    ]
   },
   {