Skills › Scientific Agent Skills › pyhealth › starter_pipeline.py
starter_pipeline.py
108 lines · skills/pyhealth/assets/starter_pipeline.py at 5f2bc03
1"""PyHealth 2.0.2 CPU starter: --demo or --root <MIMIC-III directory/URL>.
3The demo has invented codes and labels: it checks plumbing, not clinical validity.
4Real-data mode uses the upstream next-visit mortality task, not early warning.
5"""
7import argparse
8from pathlib import Path
11def build_samples(root=None, cache_dir=None, dev=False):
12 from pyhealth.datasets import MIMIC3Dataset, create_sample_dataset
13 from pyhealth.tasks import MortalityPredictionMIMIC3
15 task = MortalityPredictionMIMIC3()
16 if root is not None:
17 base = MIMIC3Dataset(
18 root=root,
19 tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
20 cache_dir=cache_dir,
21 dev=dev,
22 num_workers=1,
23 )
24 return base.set_task(task)
25 records = [
26 {
27 "patient_id": f"demo-{i}",
28 "visit_id": f"demo-{i}-{j}",
29 "conditions": ["invented-condition", f"group-{i % 3}"],
30 "procedures": ["invented-procedure"],
31 "drugs": ["invented-drug"],
32 "mortality": (i + j) % 2,
33 }
34 for i in range(60)
35 for j in range(2)
36 ]
37 return create_sample_dataset(
38 records, input_schema=task.input_schema, output_schema=task.output_schema
39 )
42def train_and_evaluate(samples, output_path, epochs=1, seed=42):
43 import random
44 import numpy as np
45 import torch
46 from pyhealth.datasets import get_dataloader, split_by_patient
47 from pyhealth.models import Transformer
48 from pyhealth.trainer import Trainer
50 if epochs < 1:
51 raise ValueError("epochs must be positive")
52 random.seed(seed)
53 np.random.seed(seed)
54 torch.manual_seed(seed)
55 torch.set_num_threads(1)
56 parts = split_by_patient(samples, [0.6, 0.2, 0.2], seed=seed)
57 patient_sets = [set(part.patient_to_index) for part in parts]
58 for i in range(3):
59 for j in range(i):
60 if patient_sets[i] & patient_sets[j]:
61 raise ValueError("Patient overlap between partitions")
62 for name, part in zip(("train", "validation", "test"), parts):
63 if len(part) == 0:
64 raise ValueError(f"Empty {name} partition")
65 labels = {int(part[i]["mortality"].item()) for i in range(len(part))}
66 if labels != {0, 1}:
67 raise ValueError(f"{name} needs both binary classes for ROC AUC")
68 loaders = [
69 get_dataloader(part, batch_size=16, shuffle=(i == 0))
70 for i, part in enumerate(parts)
71 ]
72 model = Transformer(dataset=parts[0], embedding_dim=16, dropout=0.0)
73 trainer = Trainer(
74 model=model,
75 metrics=["pr_auc", "roc_auc", "f1"],
76 device="cpu",
77 output_path=str(output_path),
78 exp_name="starter",
79 )
80 trainer.train(
81 train_dataloader=loaders[0],
82 val_dataloader=loaders[1],
83 epochs=epochs,
84 monitor="pr_auc",
85 monitor_criterion="max",
86 patience=5,
87 )
88 scores = trainer.evaluate(loaders[2])
89 return trainer, scores
92def main():
93 parser = argparse.ArgumentParser(description=__doc__)
94 source = parser.add_mutually_exclusive_group(required=True)
95 source.add_argument("--demo", action="store_true", help="Use invented in-memory data")
96 source.add_argument("--root", help="Local or public MIMIC-III root")
97 parser.add_argument("--cache-dir", type=Path, default=Path("cache/mimic3"))
98 parser.add_argument("--output", type=Path, default=Path("output/pyhealth"))
99 parser.add_argument("--epochs", type=int, default=1)
100 parser.add_argument("--dev", action="store_true", help="Limit raw data to 1000 patients")
101 args = parser.parse_args()
102 samples = build_samples(args.root, args.cache_dir, args.dev)
103 _, scores = train_and_evaluate(samples, args.output, epochs=args.epochs)
104 print("[OK]", scores)
107if __name__ == "__main__":
108 main()