cocoa-tokenizer 26.6.0__py3-none-any.whl

cocoa/cli.py

@@ -0,0 +1,285 @@
+#!/usr/bin/env python3
+
+"""
+CLI for cocoa - configurable collation and tokenization
+"""
+
+import pathlib
+import time
+from importlib.metadata import version
+from typing import Annotated, Optional
+
+import typer
+from rich import print
+from rich.console import Console
+
+from cocoa.collator import Collator
+from cocoa.tokenizer import Tokenizer
+from cocoa.util import combine_processed_data
+from cocoa.winnower import Winnower
+
+__version__ = version("cocoa")
+
+app = typer.Typer(
+    name="cocoa",
+    help=f"Configurable collation and tokenization (v{__version__})",
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+console = Console()
+
+
+@app.command()
+def collate(
+    collation_config: Annotated[
+        Optional[pathlib.Path],
+        typer.Option(
+            "--collation-config",
+            "-c",
+            help="Collation configuration file (overrides default)",
+            show_default=False,
+        ),
+    ] = None,
+    raw_data_home: Annotated[
+        str, typer.Option("--raw-data-home", "-r", help="Raw data directory")
+    ] = ...,
+    processed_data_home: Annotated[
+        str,
+        typer.Option("--processed-data-home", "-p", help="Processed data directory"),
+    ] = ...,
+    verbose: Annotated[
+        bool,
+        typer.Option(
+            "--verbose",
+            "-v",
+            help="Verbose logging for collate; this may cause "
+            "memory issues with large datasets",
+            is_flag=True,
+        ),
+    ] = False,
+):
+    """
+    Collate raw data into a denormalized format.
+
+    Reads collation configuration and produces a MEDS-like parquet file
+    with collated events.
+    """
+    with console.status("[bold green]Collating data..."):
+        t0 = time.perf_counter()
+        collator = Collator(
+            collation_cfg=collation_config,
+            raw_data_home=raw_data_home,
+            processed_data_home=processed_data_home,
+        )
+        collator.save_all(verbose=verbose)
+        t1 = time.perf_counter()
+        print(f"\n[green]✓[/green] Collation completed in {t1 - t0:.2f}s.")
+    out_path = collator.processed_data_home
+    print(f"  Output: [cyan]{out_path}/meds.parquet[/cyan]")
+    print(f"  Output: [cyan]{out_path}/subject_splits.parquet[/cyan]")
+
+
+@app.command()
+def tokenize(
+    tokenization_config: Annotated[
+        Optional[pathlib.Path],
+        typer.Option(
+            "--tokenization-config",
+            "-c",
+            help="Tokenization configuration file (overrides default)",
+            show_default=False,
+        ),
+    ] = None,
+    processed_data_home: Annotated[
+        str,
+        typer.Option("--processed-data-home", "-p", help="Processed data directory"),
+    ] = ...,
+    tokenizer_home: Annotated[
+        Optional[str],
+        typer.Option(
+            "--tokenizer-home",
+            "-t",
+            help="Use a pretrained tokenizer at this path",
+            show_default=False,
+        ),
+    ] = None,
+    verbose: Annotated[
+        bool,
+        typer.Option(
+            "--verbose",
+            "-v",
+            help="Verbose logging for collate; this may cause "
+            "memory issues with large datasets",
+            is_flag=True,
+        ),
+    ] = False,
+):
+    """
+    Tokenize collated data into integer sequences.
+
+    Reads collated parquet files and produces tokenized timelines with
+    vocabulary and bin information.
+    """
+    with console.status("[bold green]Tokenizing data..."):
+        t0 = time.perf_counter()
+        if tokenizer_home is not None:
+            print(f"Using pretrained tokenizer from [cyan]{tokenizer_home}[/cyan]...")
+            tokenizer = Tokenizer(
+                tokenization_cfg=tokenization_config,
+                processed_data_home=processed_data_home,
+            ).load(tokenizer_home)
+        else:
+            tokenizer = Tokenizer(
+                tokenization_cfg=tokenization_config,
+                processed_data_home=processed_data_home,
+            )
+        tokenizer.save_all(verbose=verbose)
+        t1 = time.perf_counter()
+        print(f"\n[green]✓[/green] Tokenization completed in {t1 - t0:.2f}s.")
+    out_path = tokenizer.processed_data_home
+    print(f"  Output: [cyan]{out_path}/tokens_times.parquet[/cyan]")
+    print(f"  Output: [cyan]{out_path}/tokenizer.yaml[/cyan]")
+    print(f"  Vocabulary size: [cyan]{len(tokenizer)}[/cyan] tokens")
+
+
+@app.command()
+def winnow(
+    winnowing_config: Annotated[
+        Optional[pathlib.Path],
+        typer.Option(
+            "--winnowing-config",
+            "-c",
+            help="Winnowing configuration file (overrides default)",
+            show_default=False,
+        ),
+    ] = None,
+    processed_data_home: Annotated[
+        str,
+        typer.Option("--processed-data-home", "-p", help="Processed data directory"),
+    ] = ...,
+    verbose: Annotated[
+        bool,
+        typer.Option(
+            "--verbose",
+            "-v",
+            help="Verbose logging for winnow; prints summary statistics",
+            is_flag=True,
+        ),
+    ] = False,
+):
+    """
+    Winnow held-out data for evaluation.
+
+    Filters held-out timelines and assigns flags to disqualify certain subjects
+    from evaluation based on the configured criteria.
+    """
+    with console.status("[bold green]Winnowing data..."):
+        t0 = time.perf_counter()
+        winnower = Winnower(
+            winnowing_cfg=winnowing_config, processed_data_home=processed_data_home
+        )
+        winnower.save_all(verbose=verbose)
+        t1 = time.perf_counter()
+        print(f"\n[green]✓[/green] Winnowing completed in {t1 - t0:.2f}s.")
+    out_path = winnower.processed_data_home
+    for s in winnower.cfg.get("splits", ["held_out"]):
+        print(f"  Output: [cyan]{out_path}/{s}_for_inference.parquet[/cyan]")
+
+
+@app.command()
+def pipeline(
+    collation_config: Annotated[
+        Optional[pathlib.Path],
+        typer.Option(
+            "--collation-config",
+            help="Collation configuration file (overrides default)",
+            show_default=False,
+        ),
+    ] = None,
+    tokenization_config: Annotated[
+        Optional[pathlib.Path],
+        typer.Option(
+            "--tokenization-config",
+            help="Tokenization configuration file (overrides default)",
+            show_default=False,
+        ),
+    ] = None,
+    winnowing_config: Annotated[
+        Optional[pathlib.Path],
+        typer.Option(
+            "--winnowing-config",
+            help="Winnowing configuration file (overrides default)",
+            show_default=False,
+        ),
+    ] = None,
+    raw_data_home: Annotated[
+        str, typer.Option("--raw-data-home", "-r", help="Raw data directory")
+    ] = ...,
+    processed_data_home: Annotated[
+        str,
+        typer.Option("--processed-data-home", "-p", help="Processed data directory"),
+    ] = ...,
+    verbose: Annotated[
+        bool,
+        typer.Option(
+            "--verbose", "-v", help="Verbose logging for pipeline steps", is_flag=True
+        ),
+    ] = False,
+):
+    """
+    Run the full pipeline: collate, tokenize, & winnow.
+    """
+    print("[bold]Running full pipeline[/bold]\n")
+    t0 = time.perf_counter()
+    collate(
+        collation_config=collation_config,
+        raw_data_home=raw_data_home,
+        processed_data_home=processed_data_home,
+        verbose=verbose,
+    )
+    tokenize(
+        tokenization_config=tokenization_config,
+        processed_data_home=processed_data_home,
+        verbose=verbose,
+    )
+    winnow(
+        winnowing_config=winnowing_config,
+        processed_data_home=processed_data_home,
+        verbose=verbose,
+    )
+    t1 = time.perf_counter()
+    print(f"\n[bold green]Pipeline completed in {t1 - t0:.2f}s.[/bold green]")
+
+
+@app.command()
+def combine_datasets(
+    input_data_dirs: list[str],
+    output_data_dir: Annotated[
+        str,
+        typer.Option(
+            "--output-data-dir", "-o", help="Output directory for combined data"
+        ),
+    ] = ...,
+):
+    """
+    Combine multiple processed datasets into one.
+
+    Merges parquet files and validates that tokenizer configurations match
+    across all input directories.
+    """
+    with console.status("[bold green]Combining datasets..."):
+        t0 = time.perf_counter()
+        output = combine_processed_data(
+            processed_data_homes=input_data_dirs, processed_data_home=output_data_dir
+        )
+        t1 = time.perf_counter()
+        print(f"\n[green]✓[/green] Combine completed in {t1 - t0:.2f}s.")
+        print(f"  Output at: [cyan]{output}[/cyan]")
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    # pipeline()
+    main()

cocoa/collator.py

@@ -0,0 +1,318 @@
+#!/usr/bin/env python3
+
+"""
+collects and collates different dataframes into a denormalized format
+"""
+
+import pathlib
+
+import numpy as np
+import polars as pl
+
+from cocoa.configurable import Configurable
+
+
+class Collator(Configurable):
+    default_file = "collation.yaml"
+
+    def __init__(
+        self,
+        collation_cfg: pathlib.Path | str = None,
+        raw_data_home: pathlib.Path | str = None,
+        processed_data_home: pathlib.Path | str = None,
+        **kwargs,
+    ):
+        super().__init__(collation_cfg, **kwargs)
+
+        self.raw_data_home, self.processed_data_home = map(
+            lambda p: pathlib.Path(p).expanduser().resolve(),
+            [raw_data_home, processed_data_home],
+        )
+        self.processed_data_home.mkdir(parents=True, exist_ok=True)
+        self.reference_frame = None
+        self.splits: tuple = ("train", "tuning", "held_out")
+        self._tz_warned: set = set()
+
+        self.logger.info("Collator initialized...")
+        self.logger.info(f"{self.raw_data_home=}")
+        self.logger.info(f"{self.processed_data_home=}")
+
+    @staticmethod
+    def slightly_safer_eval(expr):
+        """
+        performs eval on polars expressions;
+        prevents some forms of accidental usage
+        but is not secure against malicious input
+        """
+        return eval(
+            expr,
+            {"__builtins__": {"str": str, "int": int, "float": float, "bool": bool}},
+            {"pl": pl},
+        )
+
+    def to_utc_naive(
+        self, df: pl.LazyFrame, column: str, *, table: str = None
+    ) -> pl.Expr:
+        """
+        expression converting `column` to a timezone-naive UTC datetime;
+        tz-aware columns are converted to UTC (instant-preserving),
+        tz-naive columns are assumed to already be UTC
+        """
+        dtype = df.collect_schema().get(column)
+        if isinstance(dtype, pl.Datetime) and dtype.time_zone not in (None, "UTC"):
+            if (key := (table, column, dtype.time_zone)) not in self._tz_warned:
+                self._tz_warned.add(key)
+                self.logger.warning(
+                    f"{table or '<frame>'}.{column}: "
+                    f"converting {dtype.time_zone} -> UTC"
+                )
+            return (
+                pl.col(column)
+                .dt.convert_time_zone("UTC")
+                .dt.replace_time_zone(time_zone=None)
+            )
+        return pl.col(column).cast(pl.Datetime).dt.replace_time_zone(time_zone=None)
+
+    def load_table(
+        self,
+        *,
+        table: str = None,
+        filter_expr: str | list = None,
+        agg_expr: str | list = None,
+        with_col_expr: str | list = None,
+        key: str = None,
+        subject_id_str: str = None,
+        time: str = None,
+        **kwargs,
+    ) -> pl.LazyFrame:
+        """
+        lazy-load the table `table`.parquet and perform some standard ETL
+        tasks in the following order if specified:
+        1. fix subject_id to `subject_id_str`
+        2. perform a filter operation
+        3. add columns
+        4. perform an aggregation by `key` or self.cfg["subfject_id"]
+        """
+        if (f := self.raw_data_home / f"{table}.parquet").exists():
+            df = pl.scan_parquet(f)
+        elif (f := self.raw_data_home / f"{table}.csv").exists():
+            df = pl.scan_csv(f)
+        else:
+            raise FileNotFoundError(
+                f"No parquet / csv file found for {table=} in {self.raw_data_home}"
+            )
+        if subject_id_str is not None:
+            df = df.with_columns(pl.col(subject_id_str).alias(self.cfg["subject_id"]))
+        if filter_expr is not None:
+            df = df.filter(
+                self.slightly_safer_eval(filter_expr)
+                if isinstance(filter_expr, str)
+                else [self.slightly_safer_eval(c) for c in filter_expr]
+            )
+        if with_col_expr is not None:
+            df = df.with_columns(
+                self.slightly_safer_eval(with_col_expr)
+                if isinstance(with_col_expr, str)
+                else [self.slightly_safer_eval(c) for c in with_col_expr]
+            )
+        if agg_expr is not None:
+            df = (
+                df.sort(time)
+                .group_by(key if key is not None else self.cfg["subject_id"])
+                .agg(
+                    self.slightly_safer_eval(agg_expr)
+                    if isinstance(agg_expr, str)
+                    else [self.slightly_safer_eval(c) for c in agg_expr]
+                )
+            )
+        return df
+
+    def get_reference_frame(self) -> pl.LazyFrame:
+        """create the static reference frame as configured"""
+        cfg = self.cfg["reference"]
+        if self.reference_frame is not None:  # pull from cache if available
+            return self.reference_frame
+        df = self.load_table(**cfg)
+        for tkv in cfg.get("augmentation_tables", []):
+            df = df.join(
+                self.load_table(**tkv),
+                on=tkv["key"],
+                validate=tkv["validation"],
+                how="left",
+                maintain_order="left",
+            )
+        for col in (cfg["start_time"], cfg["end_time"]):
+            df = df.with_columns(
+                self.to_utc_naive(df, col, table=cfg["table"]).alias(col)
+            )
+        self.reference_frame = df  # cache result
+        return self.reference_frame
+
+    def get_entry(
+        self,
+        table: str,
+        time: str,
+        code: str,
+        *,
+        prefix: str = None,
+        numeric_value: str = None,
+        text_value: str = None,
+        filter_expr: str = None,
+        with_col_expr: str = None,
+        reference_key: str = None,
+        subject_id_str: str = None,
+        fix_date_to_time: bool = None,
+        agg_expr: str | list = None,
+        key: str = None,
+    ) -> pl.LazyFrame:
+        """create tokens corresponding to a configured event"""
+        df = (
+            self.get_reference_frame()
+            if table == "REFERENCE"
+            else self.load_table(
+                table=table,
+                time=time,
+                filter_expr=filter_expr,
+                with_col_expr=with_col_expr,
+                subject_id_str=subject_id_str,
+                agg_expr=agg_expr,
+                key=key,
+            )
+        )
+        if time in df.collect_schema().names():
+            df = df.with_columns(self.to_utc_naive(df, time, table=table).alias(time))
+        # otherwise `time` arrives normalized via the reference_key join below
+        if fix_date_to_time:
+            # if a date was cast to a time,
+            # the default of 00:00:00 should be replaced with 23:59:59
+            df = df.with_columns(
+                pl.when(pl.col(time).cast(pl.Datetime).dt.time() == pl.time(0, 0, 0))
+                .then(
+                    pl.col(time)
+                    .cast(pl.Datetime)
+                    .dt.replace(hour=23, minute=59, second=59)
+                )
+                .otherwise(pl.col(time).cast(pl.Datetime))
+                .alias(time)
+            )
+        if reference_key is not None:
+            df = df.join(self.reference_frame, on=reference_key, how="inner").filter(
+                pl.col(time)
+                .cast(pl.Datetime)
+                .is_between(
+                    pl.col(self.cfg["reference"]["start_time"]).cast(pl.Datetime),
+                    pl.col(self.cfg["reference"]["end_time"]).cast(pl.Datetime),
+                )
+            )
+
+        return df.select(
+            pl.col(self.cfg["subject_id"]).cast(pl.String).alias("subject_id"),
+            # `time` was normalized to naive UTC when the frame was loaded above
+            pl.col(time).alias("time"),
+            pl.when(pl.col(code).is_not_null())
+            .then(
+                pl.concat_str(
+                    [
+                        pl.lit(prefix),
+                        pl.col(code)
+                        .cast(pl.String)
+                        .str.to_lowercase()
+                        .str.replace_all(r"\s+", "_"),
+                    ],
+                    separator="//",
+                    ignore_nulls=True,  # prefix is optional
+                )
+            )
+            .alias("code"),
+            (pl.col(numeric_value) if numeric_value else pl.lit(None))
+            .cast(pl.Float32)  # dumb
+            .alias("numeric_value"),
+            (
+                pl.col(text_value).str.to_lowercase().str.replace_all(r"\s+", "_")
+                if text_value
+                else pl.lit(None)
+            )
+            .cast(pl.String)
+            .alias("text_value"),
+        ).drop_nulls(subset=["subject_id", "time", "code"])
+
+    def get_all(self) -> pl.LazyFrame:
+        """get all tokens for all events as configured"""
+        return pl.concat(
+            (self.get_entry(**entry) for entry in self.cfg.get("entries", []))
+        )
+
+    def get_subject_splits(self) -> pl.DataFrame:
+        """get the subject splits as configured"""
+        partition = (
+            self.get_reference_frame()
+            .group_by(self.cfg.get("group_id", self.cfg["subject_id"]))
+            .agg(pl.col(self.cfg.reference.start_time).min().alias("first_time"))
+            .sort("first_time")
+            .with_row_index()
+        ).collect()
+        split_idx = (
+            (
+                len(partition)
+                * np.cumsum(
+                    [
+                        self.cfg.subject_splits.train_frac,
+                        self.cfg.subject_splits.tuning_frac,
+                    ]
+                )
+            )
+            .astype(int)
+            .tolist()
+        )
+        partition = partition.with_columns(
+            split=pl.when(pl.col("index") < split_idx[0])
+            .then(pl.lit(self.splits[0]))
+            .when(pl.col("index") < split_idx[1])
+            .then(pl.lit(self.splits[1]))
+            .otherwise(pl.lit(self.splits[2]))
+        )
+        return (
+            partition
+            if "group_id" not in self.cfg
+            else self.get_reference_frame()
+            .collect()
+            .join(partition, on=self.cfg["group_id"])
+        ).select(pl.col(self.cfg["subject_id"]).alias("subject_id"), "split")
+
+    def save_all(self, verbose: bool = False):
+        """save collated data and subject splits to disc, optionally w/ summary stats"""
+        meds_path = self.processed_data_home / "meds.parquet"
+        meds_path.unlink(missing_ok=True)
+        df_all = self.get_all()
+        try:
+            df_all.sink_parquet(meds_path, engine="streaming")
+        except Exception as e:
+            self.logger.warning(f"Streaming write failed: {e}")
+            df_all.sink_parquet(meds_path, engine="in-memory", mkdir=True)
+        df_splits = self.get_subject_splits()
+        if "pass_through_columns" in self.cfg:
+            df_splits = df_splits.join(
+                self.reference_frame.select(
+                    pl.col(self.cfg["subject_id"]).alias("subject_id"),
+                    *self.cfg.pass_through_columns,
+                ).collect(),
+                on="subject_id",
+                how="inner",
+                validate="1:1",
+            )
+        df_splits.write_parquet(
+            self.processed_data_home / "subject_splits.parquet", mkdir=True
+        )
+
+        if verbose:
+            self.logger.summarize_meds_like(df_all, df_splits)
+
+
+if __name__ == "__main__":
+    self = Collator(
+        raw_data_home="./raw-data/raw-mimic/dev/",
+        processed_data_home="./processed/mimic/",
+    )
+    self.save_all(verbose=True)
+    print(self.get_subject_splits())
+    # breakpoint()

cocoa/config/__init__.py

@@ -0,0 +1 @@
+"""Packaged default configuration resources for cocoa."""