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.engine = create_engine(
            model_path=str(self.model_home),
            max_len=self.cfg.score.max_len,
            use_time_horizon="max_time" in self.cfg.score,  # use if max_time configured
        )

        self.ds = pl.scan_parquet(
            self.processed_data_home / "held_out_for_inference.parquet"
        )
        self.tokens_past = self.ds.select("tokens_past").collect().to_series().to_list()

        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=}"
        )

    async def sco_re(self, target_token: str, to_score_tokens: list[int]):
        tid = self.tkzr_cfg.lookup[target_token]
        sco_re_config = GenerationConfig(
            max_len=self.cfg.score.max_len,
            n_samp=self.cfg.score.n_samp,
            target_event_id=tid,
            end_token_ids=set(map(self.tkzr_cfg.lookup.get, self.cfg.score.end_tokens)),
            suppressed_ids=list(
                map(self.tkzr_cfg.lookup.get, self.cfg.score.suppressed_tokens)
            ),
            trunc_id=self.tkzr_cfg.lookup.get(self.cfg.score.trunc_id, -1),
            max_time=self.cfg.score.get("max_time", None),
        )
        trajectories, results = await generate_and_score(
            self.engine, sco_re_config, to_score_tokens, target_token_id=tid
        )
        return trajectories, results

    async def score(self):
        res = collections.defaultdict(lambda: np.nan * np.ones(len(self.tokens_past)))

        for tt in tqdm.tqdm(self.grokked_outcome_tokens, position=0):
            to_score = (
                self.ds.select(~pl.col(f"{tt}_past")).collect().to_series().to_numpy()
            )
            to_score_idx = np.flatnonzero(to_score)
            to_score_tokens = [
                x[
                    -self.cfg.score.max_len + 100 :
                ]  # allow some extra room for generation
                for x, flag in zip(self.tokens_past, to_score)
                if flag
            ]
            for idx_tks in tqdm.tqdm(
                batched(enumerate(to_score_tokens), self.cfg.score.batch_size),
                position=1,
                leave=False,
                total=np.ceil(len(to_score_tokens) / self.cfg.score.batch_size),
            ):
                idx, tks = zip(*idx_tks)
                _, results = await self.sco_re(tt, tks)
                rows = to_score_idx[np.array(idx).ravel()]
                res[f"{tt}_mc_score"][rows] = np.array(
                    [np.mean(r.m0_samples) for r in results]
                )
                res[f"{tt}_scope_score"][rows] = np.array(
                    [np.mean(r.m1_samples) for r in results]
                )
                res[f"{tt}_reach_score"][rows] = np.array(
                    [np.mean(r.m2_samples) for r in results]
                )
        return res

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

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

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
                )
            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
                )

        mdl.fit(
            X=self.features["train"][train_valid],
            y=train_label[train_valid],
            **(
                {
                    "eval_set": [
                        (
                            self.features["tuning"][tuning_valid],
                            tuning_label[tuning_valid],
                        )
                    ],
                    "eval_metric": "auc",
                }
                if str(self.estimator_type).lower() in ("lightgbm", "xgboost")
                else {}
            ),
        )

        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 score(self):
        res = dict()
        for tt in tqdm.tqdm(self.grokked_outcome_tokens, position=0):
            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)

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"