carejev-omni / example.py
nevermindai's picture
CareJev-Omni release
19d52b2 verified
Raw History Blame Contribute Delete
3.21 kB
"""Runnable example shipped inside the export (copied next to the loader shim as ``example.py``).
pip install -r requirements.txt
python example.py # uses this directory as the checkpoint
Prints one answer of each question type. Needs a GPU with room for the checkpoint; pass a directory
as the first argument to score a different export.
"""
import sys
from pathlib import Path
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from carejev_omni_loader import load_carejev_omni # noqa: E402 (shipped next to this file)
STATE = (
"PATIENT: 68-year-old male. Race: white.\n"
"STUDY ENTRY: hospital day 2; primary disease category: acute respiratory failure or multiple "
"organ system failure with sepsis (ARF/MOSF).\n"
"COMORBIDITIES: 2 recorded (diabetes).\n"
"SEVERITY: SUPPORT coma score 0; APACHE III physiology score 61; SUPPORT physiology score 33.9.\n"
"FUNCTION: 1 activities of daily living impaired (0 = independent, 7 = fully dependent).\n"
"DAY-3 PHYSIOLOGY: mean arterial pressure 72 mmHg; heart rate 104 /min; respiratory rate 28 /min; "
"temperature 38.1 °C; white cell count 15.2 x10^9/L; creatinine 2.1 mg/dL; sodium 139 mmol/L; "
"arterial pH 7.33; blood urea nitrogen 44 mg/dL; urine output 900 mL/day."
)
LOS_LEVELS = [
"Less than 1 day",
"1 day",
"2 days",
"3 days",
"4 days",
"5 days",
"6 days",
"7 days",
"8 to 14 days",
"More than 14 days",
]
def main() -> None:
path = sys.argv[1] if len(sys.argv) > 1 else str(HERE)
clf = load_carejev_omni(path)
# noul: a yes/no question. The answer is P(yes); option order never changes it.
noul = clf.predict(
state=STATE,
question="Will this patient die before being discharged from this hospitalization?",
options=["Yes", "No"],
)
print("noul P(yes) =", round(noul["probabilities"]["Yes"], 3))
# choice: one of k options. Probabilities are averaged over rotations of a canonical order.
choice = clf.predict(
state=STATE,
question="Which single problem best explains this presentation?",
options=["Sepsis", "Heart failure", "Pulmonary embolism", "Diabetic ketoacidosis"],
)
print("choice ->", choice["prediction"], "| confidence", round(choice["confidence"], 3))
# score: ordered levels. `score` is the expected level, i.e. sum(index * probability).
score = clf.predict(
state=STATE,
question="How long will this patient's hospital stay be in total?",
options=LOS_LEVELS,
)
top = max(score["probabilities"], key=score["probabilities"].get)
expected = sum(i * p for i, p in enumerate(score["probabilities"].values()))
print(f"score -> {top} | expected level {expected:.2f} of {len(LOS_LEVELS) - 1}")
# Images, heart sounds and other media go in the same call:
# clf.predict(state=..., question=..., options=[...], media="ecg.png", modality="image")
# clf.predict(state=..., question=..., options=["Yes", "No"], media="pcg.wav", modality="audio")
print("\nResearch use only. Not a medical device. See README.md.")
if __name__ == "__main__":
main()