Scoring

Scoring turns a trained model into per-timeline outcome predictions. Cotorra offers two complementary approaches, both driven by a scoring.yaml configuration and both operating on the held-out split: generation-based scoring, which lets the model simulate the future directly, and representation-based scoring, which fits a classifier on extracted features. In each case the target_tokens glob patterns in the configuration select which vocabulary tokens count as outcomes of interest.

GenerativeScorer

GenerativeScorer predicts outcomes by generating them. Using the quick_sco_re implementation of the SCORE and REACH algorithms, it autoregressively samples many possible continuations of each timeline and estimates the probability that a target outcome token occurs. For every outcome it reports three Monte-Carlo estimates — a raw occurrence score (mc), a SCOPE score, and a REACH score — computed only for subjects who have not already experienced the outcome. Generation runs asynchronously in batches and the scores are written to a parquet file.

Bases: Configurable

Source code in src/cotorra/scorer_generative.py

class GenerativeScorer(Configurable):
    default_file = "scoring.yaml"

    def __init__(
        self,
        scoring_cfg: pathlib.Path | str = None,
        processed_data_home: pathlib.Path | str = None,
        model_home: pathlib.Path | str = None,
        output_home: pathlib.Path | str = None,
        **kwargs,
    ):
        super().__init__(scoring_cfg, **kwargs)
        self.processed_data_home, self.model_home = map(
            lambda x: pathlib.Path(x).expanduser().resolve(),
            (processed_data_home, model_home),
        )
        self.output_home = (
            pathlib.Path(output_home).expanduser().resolve()
            if output_home is not None
            else self.processed_data_home
        ) / f"scores-generative-{self.model_home.name}.parquet"

        self.tkzr_cfg = OmegaConf.load(self.processed_data_home / "tokenizer.yaml")
        self.vocab: dict[str, int] = {
            str(name): int(tid) for name, tid in self.tkzr_cfg.lookup.items()
        }

        self.grokked_outcome_tokens = [
            x
            for x in self.tkzr_cfg.lookup.keys()
            if any(fnmatch.fnmatch(x, p) for p in self.cfg.score.target_tokens)
        ]
        if not self.grokked_outcome_tokens:
            raise ValueError(
                "No vocabulary tokens matched score.target_tokens="
                f"{list(self.cfg.score.target_tokens)!r}"
            )
        self.logger.info(
            f"Processed expressions to generate {self.grokked_outcome_tokens=}"
        )
        tracked_ids = [self.vocab[name] for name in self.grokked_outcome_tokens]

        self.gen_config = build_generation_config(
            self.cfg, self.vocab, self.grokked_outcome_tokens, tracked_ids
        )
        gen_cfg = self.cfg.get("generation", {})
        self.methods = gen_cfg.get("methods", ["M1", "M2"])
        engine_cfg = self.cfg.get("engine", {})
        self.chunk_size = engine_cfg.get("patient_chunk_size", 64)

        self.engine = create_engine(
            model_path=str(self.model_home),
            max_len=self.gen_config.max_len,
            use_time_horizon=self.gen_config.max_time is not None
            and self.gen_config.trunc_id is not None,
            mem_fraction=engine_cfg.get("mem_fraction", 0.85),
        )

        self.ds = pl.read_parquet(
            self.processed_data_home / "held_out_for_inference.parquet"
        )
        subject_ids = self.ds.select("subject_id").to_series().to_list()
        self.tokens_past = self.ds.select("tokens_past").to_series().to_list()

        self.overflow = self.cfg.get("prompt_overflow", "truncate_left")
        max_len = self.gen_config.max_len
        self.final_tokens = []
        self.final_ids = []
        self.keep_mask = []
        n_dropped = 0
        n_truncated = 0
        for sid, toks in zip(subject_ids, self.tokens_past):
            if len(toks) > max_len:
                if self.overflow == "drop":
                    n_dropped += 1
                    self.keep_mask.append(False)
                    continue
                elif self.overflow == "truncate_left":
                    toks = toks[-max_len:]
                    n_truncated += 1
            self.final_tokens.append(toks)
            self.final_ids.append(sid)
            self.keep_mask.append(True)
        if n_dropped:
            self.logger.info(
                f"Dropped {n_dropped} patients with prompts > {max_len} tokens"
            )
        if n_truncated:
            self.logger.info(
                f"Left-truncated {n_truncated} prompts to {max_len} tokens"
            )

        kept = np.array(self.keep_mask)
        self.outcome_past_masks: list[np.ndarray] = []
        for evt_name in self.grokked_outcome_tokens:
            past_col = f"{evt_name}_past"
            if past_col in self.ds.columns:
                mask = self.ds[past_col].to_numpy().astype(bool)[kept]
            else:
                mask = np.zeros(len(self.final_ids), dtype=bool)
            self.outcome_past_masks.append(mask)

    async def sco_re(self) -> list[list[PatientResults]]:
        """Run one inline generation pass over the whole cohort, tracking every
        outcome in `self.grokked_outcome_tokens` simultaneously."""
        outcome_ids = self.gen_config.tracked_ids
        outcome_configs = [
            dataclasses.replace(self.gen_config, target_event_id=evt_id)
            for evt_id in outcome_ids
        ]
        n_patients = len(self.final_ids)
        n_outcomes = len(self.grokked_outcome_tokens)
        all_results: list[list[PatientResults]] = [[] for _ in range(n_outcomes)]
        per_outcome_m2_tokens: list[int] = [0] * n_outcomes
        per_outcome_m2_count: list[int] = [0] * n_outcomes

        with tqdm.tqdm(total=n_patients, desc="Generating") as pbar:
            for chunk_start in range(0, n_patients, self.chunk_size):
                chunk_end = min(chunk_start + self.chunk_size, n_patients)
                chunk_tokens = self.final_tokens[chunk_start:chunk_end]

                # Generate the initial M1 chunk
                chunk_m1 = await generate_trajectories(
                    self.engine,
                    self.gen_config,
                    chunk_tokens,
                    ["M1"],
                    stop_at_tracked_events=False,
                )

                outcome_m1_results_list = [
                    aggregate_inline_results(chunk_m1, len(chunk_tokens), oc)
                    for oc in outcome_configs
                ]
                for k, om1r in enumerate(outcome_m1_results_list):
                    all_results[k].extend(om1r)

                # M2 regen within this batch — avoids accumulating all regen
                # requests and triggering them in one giant post-loop gather.
                if "M2" in self.methods:
                    chunk_regen: list[tuple] = []
                    for traj in chunk_m1:
                        global_idx = chunk_start + traj.patient_idx
                        if (
                            traj.inline_tracked_ids is None
                            or traj.reach_estimates is None
                        ):
                            continue
                        for k, (evt_id, oc) in enumerate(
                            zip(outcome_ids, outcome_configs)
                        ):
                            # skip regenning this outcome if it occurred in the prefix
                            if self.outcome_past_masks[k][global_idx]:
                                continue
                            try:
                                ki = traj.inline_tracked_ids.index(evt_id)
                            except ValueError:
                                continue
                            # Prefer occurred_flag (O(1), correct for time-truncated
                            # trajectories where output_ids is already trimmed).
                            if traj.occurred_flag is not None:
                                evt_occurred = bool(traj.occurred_flag[ki])
                            else:
                                evt_occurred = evt_id in traj.output_ids
                            if evt_occurred:
                                chunk_regen.append((traj, global_idx, k, oc))
                            else:
                                all_results[k][global_idx].m2_samples.append(
                                    float(traj.reach_estimates[ki])
                                )

                    if chunk_regen:
                        m2_trajs = list(
                            await asyncio.gather(
                                *[
                                    generate_m2_from_m1_trajectory(
                                        self.engine,
                                        oc,
                                        traj,
                                        self.final_tokens[global_idx],
                                    )
                                    for traj, global_idx, _, oc in chunk_regen
                                ]
                            )
                        )
                        scored_m2 = list(
                            await asyncio.gather(
                                *[
                                    score_trajectory(
                                        self.engine,
                                        oc,
                                        m2_traj,
                                        self.final_tokens[global_idx],
                                    )
                                    for (_, global_idx, _, oc), m2_traj in zip(
                                        chunk_regen, m2_trajs
                                    )
                                ]
                            )
                        )
                        for (traj, global_idx, k, oc), st in zip(
                            chunk_regen, scored_m2
                        ):
                            all_results[k][global_idx].m2_samples.append(st.score)
                            per_outcome_m2_tokens[k] += st.trajectory.n_new_tokens or 0
                            per_outcome_m2_count[k] += 1

                pbar.update(len(chunk_tokens))

        return all_results

    async def score(self):
        all_results = await self.sco_re()
        orig_idx = np.flatnonzero(self.keep_mask)
        n = len(self.tokens_past)
        res = {}

        for k, tt in enumerate(tqdm.tqdm(self.grokked_outcome_tokens)):
            results_k = all_results[k]
            # Patients where this outcome already occurred in the past are
            # excluded from that outcome's scores (NaN).
            for i, is_past in enumerate(self.outcome_past_masks[k]):
                if is_past:
                    results_k[i] = PatientResults()

            m0 = np.nan * np.ones(n)
            m1 = np.nan * np.ones(n)
            m2 = np.nan * np.ones(n)
            m0[orig_idx] = [
                np.mean(r.m0_samples) if r.m0_samples else np.nan for r in results_k
            ]
            m1[orig_idx] = [
                np.mean(r.m1_samples) if r.m1_samples else np.nan for r in results_k
            ]
            m2[orig_idx] = [
                np.mean(r.m2_samples) if r.m2_samples else np.nan for r in results_k
            ]

            res[f"{tt}_mc_score"] = m0
            res[f"{tt}_scope_score"] = m1
            res[f"{tt}_reach_score"] = m2

        return res

    def save_all(self, verbose: bool = False):
        res = asyncio.run(self.score())
        (df_res := self.ds.with_columns(pl.from_dict(res))).write_parquet(
            self.output_home
        )

        if verbose:
            self.logger.summarize_preds(df_res, self.grokked_outcome_tokens)

sco_re() async

Run one inline generation pass over the whole cohort, tracking every outcome in self.grokked_outcome_tokens simultaneously.

Source code in src/cotorra/scorer_generative.py

async def sco_re(self) -> list[list[PatientResults]]:
    """Run one inline generation pass over the whole cohort, tracking every
    outcome in `self.grokked_outcome_tokens` simultaneously."""
    outcome_ids = self.gen_config.tracked_ids
    outcome_configs = [
        dataclasses.replace(self.gen_config, target_event_id=evt_id)
        for evt_id in outcome_ids
    ]
    n_patients = len(self.final_ids)
    n_outcomes = len(self.grokked_outcome_tokens)
    all_results: list[list[PatientResults]] = [[] for _ in range(n_outcomes)]
    per_outcome_m2_tokens: list[int] = [0] * n_outcomes
    per_outcome_m2_count: list[int] = [0] * n_outcomes

    with tqdm.tqdm(total=n_patients, desc="Generating") as pbar:
        for chunk_start in range(0, n_patients, self.chunk_size):
            chunk_end = min(chunk_start + self.chunk_size, n_patients)
            chunk_tokens = self.final_tokens[chunk_start:chunk_end]

            # Generate the initial M1 chunk
            chunk_m1 = await generate_trajectories(
                self.engine,
                self.gen_config,
                chunk_tokens,
                ["M1"],
                stop_at_tracked_events=False,
            )

            outcome_m1_results_list = [
                aggregate_inline_results(chunk_m1, len(chunk_tokens), oc)
                for oc in outcome_configs
            ]
            for k, om1r in enumerate(outcome_m1_results_list):
                all_results[k].extend(om1r)

            # M2 regen within this batch — avoids accumulating all regen
            # requests and triggering them in one giant post-loop gather.
            if "M2" in self.methods:
                chunk_regen: list[tuple] = []
                for traj in chunk_m1:
                    global_idx = chunk_start + traj.patient_idx
                    if (
                        traj.inline_tracked_ids is None
                        or traj.reach_estimates is None
                    ):
                        continue
                    for k, (evt_id, oc) in enumerate(
                        zip(outcome_ids, outcome_configs)
                    ):
                        # skip regenning this outcome if it occurred in the prefix
                        if self.outcome_past_masks[k][global_idx]:
                            continue
                        try:
                            ki = traj.inline_tracked_ids.index(evt_id)
                        except ValueError:
                            continue
                        # Prefer occurred_flag (O(1), correct for time-truncated
                        # trajectories where output_ids is already trimmed).
                        if traj.occurred_flag is not None:
                            evt_occurred = bool(traj.occurred_flag[ki])
                        else:
                            evt_occurred = evt_id in traj.output_ids
                        if evt_occurred:
                            chunk_regen.append((traj, global_idx, k, oc))
                        else:
                            all_results[k][global_idx].m2_samples.append(
                                float(traj.reach_estimates[ki])
                            )

                if chunk_regen:
                    m2_trajs = list(
                        await asyncio.gather(
                            *[
                                generate_m2_from_m1_trajectory(
                                    self.engine,
                                    oc,
                                    traj,
                                    self.final_tokens[global_idx],
                                )
                                for traj, global_idx, _, oc in chunk_regen
                            ]
                        )
                    )
                    scored_m2 = list(
                        await asyncio.gather(
                            *[
                                score_trajectory(
                                    self.engine,
                                    oc,
                                    m2_traj,
                                    self.final_tokens[global_idx],
                                )
                                for (_, global_idx, _, oc), m2_traj in zip(
                                    chunk_regen, m2_trajs
                                )
                            ]
                        )
                    )
                    for (traj, global_idx, k, oc), st in zip(
                        chunk_regen, scored_m2
                    ):
                        all_results[k][global_idx].m2_samples.append(st.score)
                        per_outcome_m2_tokens[k] += st.trajectory.n_new_tokens or 0
                        per_outcome_m2_count[k] += 1

            pbar.update(len(chunk_tokens))

    return all_results

RepBasedScorer

RepBasedScorer predicts outcomes from the representations dumped by the Extractor. It loads the extracted feature vectors for the train, tuning, and held-out splits, fits a supervised classifier per outcome token to predict whether that outcome occurs, and writes the held-out predicted probabilities to a parquet file. The classifier family is chosen with EstimatorType; it errors with a helpful message if the features are missing, prompting you to run cotorra extract first.

Bases: Configurable

Source code in src/cotorra/scorer_rep_based.py

class RepBasedScorer(Configurable):
    default_file = "scoring.yaml"

    def __init__(
        self,
        scoring_cfg: pathlib.Path | str = None,
        processed_data_home: pathlib.Path | str = None,
        model_home: pathlib.Path | str = None,
        output_home: pathlib.Path | str = None,
        estimator_type: typing.Literal[
            "k-NN",
            "lightGBM",
            "logistic",
            "logistic-z",
            "logistic-CV",
            "logistic-CV-z",
            "XGBoost",
        ] = "lightGBM",
        **kwargs,
    ):
        super().__init__(scoring_cfg, **kwargs)
        self.processed_data_home, self.model_home = map(
            lambda x: pathlib.Path(x).expanduser().resolve(),
            (processed_data_home, model_home),
        )
        self.output_home = (
            pathlib.Path(output_home).expanduser().resolve()
            if output_home is not None
            else self.processed_data_home
        ) / f"scores-rep-based-{self.model_home.name}.parquet"
        self.tkzr_cfg = OmegaConf.load(self.processed_data_home / "tokenizer.yaml")

        self.splits = ("train", "tuning", "held_out")
        self.estimator_type = estimator_type

        try:
            self.features = {
                s: np.vstack(
                    pl.scan_parquet(
                        self.processed_data_home
                        / f"features-{s}*-{self.model_home.name}.parquet"
                    )
                    .select("features")
                    .collect()
                    .to_series()
                    .to_list()
                )
                for s in self.splits
            }
        except FileNotFoundError as e:
            raise FileNotFoundError(
                "Expected extracted features at: "
                f"{self.processed_data_home / 'features-<split>-<model_name>.parquet'},"
                " but not found."
                " Please run `cotorra extract` first."
            ) from e

        self.labels = {
            s: pl.scan_parquet(self.processed_data_home / f"{s}_for_inference.parquet")
            for s in self.splits
        }

        self.grokked_outcome_tokens = [
            x
            for x in self.tkzr_cfg.lookup.keys()
            if any(fnmatch.fnmatch(x, p) for p in self.cfg.score.target_tokens)
        ]
        self.logger.info(
            f"Processed expressions to generate {self.grokked_outcome_tokens=}"
        )

    def score_label(self, target_token="DSCG//expired"):
        cols = (~pl.col(f"{target_token}_past"), f"{target_token}_future")
        train_valid, train_label = (
            self.labels["train"].select(*cols).collect().to_numpy().T
        )
        tuning_valid, tuning_label = (
            self.labels["tuning"].select(*cols).collect().to_numpy().T
        )
        held_out_valid = (
            self.labels["held_out"].select(cols[0]).collect().to_numpy().ravel()
        )

        match str(self.estimator_type).lower():
            case "logistic" | "lr" | "logistic-regression":
                self.logger.info("Using logistic regression classifier")
                mdl = skl.linear_model.LogisticRegression(max_iter=10_000)
            case "logistic-z" | "lr-z" | "logistic-regression-z":
                self.logger.info(
                    "Using logistic regression classifier on z-scored features"
                )
                mdl = skl.pipeline.make_pipeline(
                    skl.preprocessing.StandardScaler(),
                    skl.linear_model.LogisticRegression(max_iter=10_000),
                )
            case "logistic-cv" | "lr-cv":
                self.logger.info(
                    "Using logistic regression classifier with cross-validation"
                )
                mdl = skl.linear_model.LogisticRegressionCV(
                    n_jobs=-1,
                    scoring="roc_auc",
                    max_iter=10_000,
                    use_legacy_attributes=False,
                    l1_ratios=(0,),
                )
            case "logistic-cv-z" | "lr-cv-z":
                self.logger.info(
                    "Using logistic regression classifier with cross-validation "
                    "on z-scored features"
                )
                mdl = skl.pipeline.make_pipeline(
                    skl.preprocessing.StandardScaler(),
                    skl.linear_model.LogisticRegressionCV(
                        n_jobs=-1,
                        scoring="roc_auc",
                        max_iter=10_000,
                        use_legacy_attributes=False,
                        l1_ratios=(0,),  # suppresses a warning
                    ),
                )
            case "k-nn" | "knn" | "k_nn":
                self.logger.info("Using k-nn classifier")
                mdl = skl.neighbors.KNeighborsClassifier(
                    n_neighbors=max(25, int(0.2 * sum(train_valid))), n_jobs=-1
                )
            case "xgboost":
                self.logger.info("Using XGBoost classifier")
                mdl = xgb.XGBClassifier(
                    min_child_weight=5,
                    max_leaves=64,
                    n_estimators=250,
                    n_jobs=-1,
                    # xgboost >= 2.0 takes `eval_metric` here rather than on
                    # `fit`, where lightGBM still wants it
                    eval_metric="auc",
                )
            case _:
                self.logger.info("Using (default) lightGBM classifier")
                mdl = lgb.LGBMClassifier(
                    min_data_in_leaf=5, num_leaves=64, n_estimators=250, n_jobs=-1
                )

        fit_kwargs = dict()
        if (estimator := str(self.estimator_type).lower()) in EVAL_SET_ESTIMATORS:
            fit_kwargs["eval_set"] = [
                (self.features["tuning"][tuning_valid], tuning_label[tuning_valid])
            ]
            if estimator == "lightgbm":
                fit_kwargs["eval_metric"] = "auc"

        mdl.fit(
            X=self.features["train"][train_valid],
            y=train_label[train_valid],
            **fit_kwargs,
        )

        scores = np.nan * np.ones_like(held_out_valid)
        scores[held_out_valid] = mdl.predict_proba(
            X=self.features["held_out"][held_out_valid]
        )[:, 1]

        return scores

    def unfittable_reason(self, target_token: str) -> str | None:
        """
        why `target_token` cannot be fit, or `None` if it can. The winnowed
        inference tables routinely hold labels no estimator can be trained
        on -- a token that never made it into the tables at all, or one whose
        rows not already past the threshold are all a single class -- so
        `score` checks before fitting rather than letting one bad label abort
        the whole run and lose the scores for every other one.
        """
        splits = ("train", "tuning")
        if str(self.estimator_type).lower() not in EVAL_SET_ESTIMATORS:
            splits = ("train",)

        for split in splits:
            cols = self.labels[split].collect_schema().names()
            if missing := [
                c
                for c in (f"{target_token}_past", f"{target_token}_future")
                if c not in cols
            ]:
                return f"{split} is missing {', '.join(missing)}"

            valid, label = (
                self.labels[split]
                .select(~pl.col(f"{target_token}_past"), f"{target_token}_future")
                .collect()
                .to_numpy()
                .T
            )
            n_classes = len(np.unique(label[valid]))
            if n_classes < 2:
                return (
                    f"{split} has {int(valid.sum())} row(s) not already past the "
                    f"threshold, covering {n_classes} class(es)"
                )

        return None

    def score(self):
        res = dict()
        for tt in tqdm.tqdm(self.grokked_outcome_tokens, position=0):
            if (reason := self.unfittable_reason(tt)) is not None:
                self.logger.warning(f"Skipping {tt}: {reason}")
                continue
            res[f"{tt}_rep_score"] = self.score_label(target_token=tt)

        return res

    def save_all(self, verbose: bool = False):
        (
            df_res := self.labels["held_out"].with_columns(pl.from_dict(self.score()))
        ).sink_parquet(self.output_home)

        if verbose:
            self.logger.summarize_preds(df_res, self.grokked_outcome_tokens)

unfittable_reason(target_token)

why target_token cannot be fit, or None if it can. The winnowed inference tables routinely hold labels no estimator can be trained on -- a token that never made it into the tables at all, or one whose rows not already past the threshold are all a single class -- so score checks before fitting rather than letting one bad label abort the whole run and lose the scores for every other one.

Source code in src/cotorra/scorer_rep_based.py

def unfittable_reason(self, target_token: str) -> str | None:
    """
    why `target_token` cannot be fit, or `None` if it can. The winnowed
    inference tables routinely hold labels no estimator can be trained
    on -- a token that never made it into the tables at all, or one whose
    rows not already past the threshold are all a single class -- so
    `score` checks before fitting rather than letting one bad label abort
    the whole run and lose the scores for every other one.
    """
    splits = ("train", "tuning")
    if str(self.estimator_type).lower() not in EVAL_SET_ESTIMATORS:
        splits = ("train",)

    for split in splits:
        cols = self.labels[split].collect_schema().names()
        if missing := [
            c
            for c in (f"{target_token}_past", f"{target_token}_future")
            if c not in cols
        ]:
            return f"{split} is missing {', '.join(missing)}"

        valid, label = (
            self.labels[split]
            .select(~pl.col(f"{target_token}_past"), f"{target_token}_future")
            .collect()
            .to_numpy()
            .T
        )
        n_classes = len(np.unique(label[valid]))
        if n_classes < 2:
            return (
                f"{split} has {int(valid.sum())} row(s) not already past the "
                f"threshold, covering {n_classes} class(es)"
            )

    return None

EstimatorType

EstimatorType enumerates the classifier families available to RepBasedScorer: k-nearest-neighbors, LightGBM (the default), XGBoost, and several logistic-regression variants (plain, standardized/z-scored, and cross-validated). It exists so the choice of estimator can be passed as a plain string on the CLI or in configuration.

Bases: str, Enum

Source code in src/cotorra/scorer_rep_based.py

class EstimatorType(str, enum.Enum):
    knn = "k-NN"
    lightgbm = "lightGBM"
    logistic = "logistic"
    logistic_z = "logistic-z"
    logistic_cv = "logistic-CV"
    logistic_cv_z = "logistic-CV-z"
    xgboost = "XGBoost"