Extraction

The extraction stage runs a trained model over timelines and captures the representations it computes, so that a lightweight classifier can later be fit on them (see RepBasedScorer). This is the representation-based counterpart to autoregressive generative scoring.

Extractor

Extractor loads a saved model, moves it to the best available device (CUDA, MPS, or CPU), and runs it in inference mode over the held-out timelines. For each timeline it reads out the model's last hidden-layer activations — by default the single vector at the final token before the first EOS (the model's summary of the timeline up to the prediction point), or, when all_times is set, the full sequence of per-token representations padded to max_seq_len. Results are written to output_home as sharded parquet files, one column of feature vectors per timeline, ready to be consumed downstream.

Bases: Configurable

load a model and extract representations from it

Source code in src/cotorra/extractor.py

class Extractor(Configurable):
    """load a model and extract representations from it"""

    default_file = "extraction.yaml"

    def __init__(
        self,
        extraction_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__(extraction_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
        )
        self.tkzr_cfg = OmegaConf.load(self.processed_data_home / "tokenizer.yaml")
        self.loader = Loader(extraction_cfg, self.processed_data_home)
        self.device = (
            "cuda"
            if t.cuda.is_available()
            else "mps"
            if t.backends.mps.is_available()
            else "cpu"
        )
        self.model = AutoModelForCausalLM.from_pretrained(self.model_home)
        self.model.to(self.device).eval()
        if not isinstance(self.model.config.pad_token_id, int):
            self.model.config.pad_token_id = self.model.config.eos_token_id
        self.ds = None

    def collate_fn(self, batch):
        ml = t.tensor(self.cfg.get("extract", {}).get("max_len", 4096))
        input_ids = pad_sequence(
            [x[:ml] for x in batch["input_ids"]],
            batch_first=True,
            padding_value=self.model.config.pad_token_id,
        ).to(self.model.device)
        if "time_based_rope" in self.cfg:
            p_ids = (
                pad_sequence(
                    [x[:ml] for x in batch["s_elapsed_past"]],
                    batch_first=True,
                    padding_value=self.model.config.pad_token_id,
                ).to(self.model.device)
                / self.cfg.time_based_rope.sec_per_pos_id
            )
            p_ids += t.arange(p_ids.shape[-1], device=p_ids.device, dtype=p_ids.dtype)
        else:
            p_ids = None
        return {"input_ids": input_ids, "position_ids": p_ids}

    def extract_final(self, batch, all_times: bool = False):
        collated = self.collate_fn(batch)
        first_eos = t.where(
            (hits := (collated["input_ids"] == self.model.config.eos_token_id)).any(
                dim=-1
            ),
            hits.long().argmax(dim=-1)
            - 1,  # -1 to get the last token before break point
            collated["input_ids"].shape[-1] - 1,
        )
        with t.inference_mode():
            features = self.model(**collated, output_hidden_states=True).hidden_states[
                -1
            ]  # last hidden layer
        if all_times:
            features = features.half().cpu().numpy()
            collated = np.full(
                shape=(features.shape[0], self.cfg.max_seq_len, features.shape[-1]),
                fill_value=np.nan,
            )
            lengths = first_eos.cpu().numpy()[:, None]
            out_mask = np.arange(collated.shape[1]) <= lengths
            feat_mask = np.arange(features.shape[1]) <= lengths
            collated[out_mask] = features[feat_mask]
            batch["features"] = collated
        else:
            batch["features"] = (
                features[t.arange(len(first_eos)), first_eos].half().cpu().numpy()
            )
        return batch

    def extract(self, all_times: bool = False):
        a = "-all" if all_times else ""
        shard_size = self.cfg.get("extract", {}).get("shard_size", None)
        ds = self.loader.for_inference.with_format("torch")
        for split, dset in ds.items():
            n = math.ceil(len(dset) / shard_size) if shard_size else 1
            for i in range(n):
                index = f"-{i:05d}-of-{n:05d}" if n > 1 else ""
                dset.shard(num_shards=n, index=i).map(
                    lambda batch: self.extract_final(batch, all_times=all_times),
                    batched=True,
                    batch_size=self.cfg.get("extract", {}).get("batch_size", 8),
                    load_from_cache_file=False,  # disable caching
                ).to_parquet(
                    self.output_home
                    / f"features{a}-{split}{index}-{self.model_home.name}.parquet"
                )