Update pinnacle.py

mims-harvard · Jan 13, 2025 · e8a5556 · e8a5556
1 parent f0ea40a
commit e8a5556
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tdc/resource/pinnacle.py b/tdc/resource/pinnacle.py
@@ -79,6 +79,8 @@ def get_embeds(self):
     def get_exp_data(self, seed=1, split="train"):
         if split not in ["train", "val", "test"]:
             raise ValueError("{} not a valid split".format(split))
+        if seed < 1 or seed > 10:
+            raise ValueError(f'{seed} is not a valid seed in range 1-10.')
         filename = "pinnacle_output{}".format(seed)
         # clean data directory
         file_list = os.listdir("./data")
@@ -121,4 +123,4 @@ def get_exp_data(self, seed=1, split="train"):
             df["disease"] = disease
             x.append(df)
             os.remove("./data/{}".format(file))
-        return pd.concat(x, axis=0, ignore_index=True)
+        return pd.concat(x, axis=0, ignore_index=True)