Tokenizer

The Tokenizer is the second stage of the pipeline. It takes the collated events and turns them into per-subject integer sequences ("timelines"), learning a vocabulary and numeric-binning scheme along the way. Its behavior is driven by a tokenization config (tokenization.yaml).

Crucially, both the vocabulary and the numeric bins are learned only on the training split, so no information leaks from tuning or held-out subjects.

What it produces

Running the tokenizer (via save_all) writes two files to the processed-data directory:

• tokens_times.parquet — the tokenized timelines (tokens, times, and optionally numeric_values) for each subject.
• tokenizer.yaml — the learned state (lookup table, bins, and config), sufficient to reconstruct the tokenizer.

How it works

The orchestrator get_all runs these steps in order:

1. add_ends — insert BOS / EOS markers at the start and end of each subject's timeline.
2. add_clocks — optionally insert CLCK//HH tokens at configured hours of the day.
3. bin_data — discretize numeric values into quantile bins (Q0, Q1, …) using cut points learned by get_bins.
4. insert_time_spacers — optionally insert TIME//… tokens encoding the gap between consecutive events.
5. tokenize_data — map the vocabulary to integers via the lookup table from get_lookup, sorting simultaneous events by the configured ordering priority. UNK is always token 0.

Reusing a trained tokenizer

A tokenizer's learned state round-trips through to_yaml and from_yaml, so a tokenizer trained on one dataset can be frozen (is_training=False) and applied to another — see the Tokenizer Transfer recipe.

!!! tip "It behaves like a mapping" A Tokenizer is callable and dict-like: tkzr("EOS") returns a token id (0 for out-of-vocabulary words), "foo" in tkzr tests vocabulary membership, and len(tkzr) reports the vocabulary size.

tokenizes collated data into integer sequences, creating bins & a lookup table

Tokenizer

Bases: Configurable

converts collated data to tokenized timelines, learning bins and lookup table on training data

Source code in src/cocoa/tokenizer.py

class Tokenizer(Configurable):
    """
    converts collated data to tokenized timelines,
    learning bins and lookup table on training data
    """

    default_file = "tokenization.yaml"

    def __init__(
        self,
        tokenization_cfg: pathlib.Path | str = None,
        processed_data_home: pathlib.Path | str = None,
        is_training: bool = True,
        **kwargs,
    ):
        super().__init__(tokenization_cfg, **kwargs)
        self.processed_data_home = (
            pathlib.Path(processed_data_home).expanduser().resolve()
        )

        self.bins = None
        self.subject_splits = None
        self.lookup = None
        self.is_training = is_training
        self.created_dttm = (
            datetime.datetime.now(zoneinfo.ZoneInfo("America/Chicago"))
            .replace(microsecond=0)
            .isoformat()
        )

    def get_data(self) -> pl.LazyFrame:
        self.logger.info(f"Loading collated data with {self.processed_data_home=}")
        self.subject_splits = pl.scan_parquet(
            self.processed_data_home / "subject_splits.parquet"
        )  # allows processed_data_home to be set after initialization
        return pl.scan_parquet(self.processed_data_home / "meds.parquet")

    @staticmethod
    def add_ends(df: pl.LazyFrame) -> pl.LazyFrame:
        """add BOS / EOS codes at appropriate places"""
        return pl.concat(
            [
                df,
                df.group_by("subject_id")
                .agg(pl.col("time").min())
                .with_columns(code=pl.lit("BOS")),
                df.group_by("subject_id")
                .agg(pl.col("time").max())
                .with_columns(code=pl.lit("EOS")),
            ],
            how="diagonal",
        )

    def add_clocks(self, df: pl.LazyFrame) -> pl.LazyFrame:
        """add clock codes if configured"""
        if self.cfg.get("insert_clocks", False):
            return pl.concat(
                [
                    df,
                    df.group_by("subject_id")
                    .agg(
                        start=pl.col("time").min().dt.truncate("1h"),
                        end=pl.col("time").max().dt.truncate("1h"),
                    )
                    .with_columns(
                        pl.datetime_ranges(
                            pl.col("start"),
                            pl.col("end"),
                            interval="1h",
                            closed="right",
                        )
                        .alias("time")
                        .list.filter(
                            pl.element().dt.strftime("%H").is_in(list(self.cfg.clocks))
                        )
                    )
                    .explode("time", keep_nulls=False)
                    .drop_nulls(subset=["time"])
                    .with_columns(HH=pl.col("time").dt.strftime("%H"))
                    .select(
                        "subject_id",
                        "time",
                        pl.concat_str(pl.lit("CLCK//"), pl.col("HH")).alias("code"),
                    ),
                ],
                how="diagonal",
            )
        else:
            return df

    def get_bins(self, df: pl.LazyFrame) -> pl.DataFrame:
        """calculate bins for numeric values on training data"""
        if self.bins is None:
            assert self.is_training, "Bins must be learned during training"
            self.bins = (
                df.join(  # restrict to training data
                    self.subject_splits.filter(pl.col("split") == "train"),
                    on="subject_id",
                    validate="m:1",
                )
                .drop_nulls(subset=["numeric_value"])
                .drop_nans(subset=["numeric_value"])
                .group_by("code")
                .agg(
                    [
                        pl.col("numeric_value")
                        .quantile(i / self.cfg.n_bins)
                        .alias(f"break_{i}")
                        for i in range(1, self.cfg.n_bins)
                    ]
                )
            ).collect()
        return self.bins

    def bin_data(self, df: pl.LazyFrame) -> pl.LazyFrame:
        """discretize numeric values with learned cut points"""
        return (
            df.join(self.get_bins(df).lazy(), on="code", how="left")
            .with_columns(
                pl.when(pl.col("numeric_value").is_not_null())
                .then(
                    pl.concat_str(
                        pl.lit("Q"),
                        pl.sum_horizontal(
                            [
                                pl.col(f"break_{i}") <= pl.col("numeric_value")
                                for i in range(1, self.cfg.n_bins)
                            ]
                        ).cast(pl.String),
                    )
                )
                .otherwise(None)
                .alias("binned_value")
            )
            .drop([f"break_{i}" for i in range(1, self.cfg.n_bins)])
        )

    def insert_time_spacers(self, df: pl.LazyFrame) -> pl.LazyFrame:
        """
        add time spacing tokens if configured;
        this should be done after clock tokens are inserted
        """
        if self.cfg.get("insert_spacers", False):
            spcrs = dict(self.cfg.spacers)
            return (
                df.sort("subject_id", "time")
                .with_columns(
                    tdiff_mins=(
                        pl.col("time") - pl.col("time").shift(1).over("subject_id")
                    ).dt.total_minutes()
                )
                .with_columns(
                    t_spacer=pl.when(
                        (pl.col("tdiff_mins") < min(spcrs.values()))
                        | pl.col("tdiff_mins").is_null()
                    )
                    .then(None)
                    .otherwise(
                        pl.concat_str(
                            pl.lit("TIME//"),
                            pl.col("tdiff_mins")
                            .cut(
                                breaks=list(spcrs.values())[1:],
                                labels=list(spcrs.keys()),
                            )
                            .cast(pl.String),
                        )
                    )
                    .cast(pl.String)
                )
            )
        else:
            return df.with_columns(t_spacer=pl.lit(None))

    def get_pretokenized(self, df: pl.LazyFrame) -> pl.LazyFrame:
        """prepare codes for tokenization, depending on whether fusion is configured"""
        return df.with_columns(
            pl.concat_list(
                pl.col("t_spacer"),
                pl.concat_str(
                    pl.col("code"),
                    pl.col("binned_value"),
                    pl.col("text_value"),
                    separator="_",
                    ignore_nulls=True,
                )
                if self.cfg.get("fused", False)
                else pl.concat_list("code", "binned_value", "text_value"),
            )
            .list.drop_nulls()
            .alias("to_tokenize")
        )

    def get_lookup(self, pt: pl.LazyFrame) -> pl.DataFrame:
        """create mapping from vocabulary to integer tokens on training data"""
        if self.lookup is None:
            assert self.is_training, "Lookup table must be learned during training"
            lookup = (
                pt.join(  # restrict to training data
                    self.subject_splits.filter(pl.col("split") == "train"),
                    on="subject_id",
                    validate="m:1",
                )
                .explode("to_tokenize")
                .select(pl.col("to_tokenize").unique().sort())
                .filter(pl.col("to_tokenize") != "UNK")  # UNK is 0
                .with_row_index("token", offset=1)
                .select("to_tokenize", "token")
            )
            unk_row = pl.LazyFrame(
                {"to_tokenize": ["UNK"], "token": pl.Series([0], dtype=pl.UInt32)}
            )
            self.lookup = pl.concat([unk_row, lookup]).collect()
        return self.lookup

    def get_priority(self) -> pl.DataFrame:
        """fetch ordering for sorting cotemporaneous codes"""
        return (
            pl.Series("code_type", self.cfg.ordering)
            .to_frame()
            .with_row_index("priority")
        )

    def tokenize_data(self, pt: pl.LazyFrame) -> pl.LazyFrame:
        """apply lookup table to pretokenized data"""
        return (
            pt.with_columns(code_type=pl.col("code").str.split("//").list[0])
            .join(
                self.get_priority().lazy(), on="code_type", how="left", validate="m:1"
            )
            .with_columns(pl.col("priority").fill_null(len(self.cfg.ordering)))
            .sort("time", "priority")
            .explode("to_tokenize")
            .join(
                self.get_lookup(pt).lazy(), on="to_tokenize", validate="m:1", how="left"
            )
            .with_columns(
                pl.col("token").fill_null(pl.lit(0, dtype=pl.UInt32))
            )  # UNK is 0
            .group_by("subject_id", maintain_order=True)
            .agg(
                pl.col("token").alias("tokens"),
                pl.col("time").alias("times"),
                *(
                    [pl.col("numeric_value").alias("numeric_values")]
                    if self.cfg.get("include_numeric_values", False)
                    else []
                ),
            )
        )

    def get_all(self, verbose: bool = False) -> pl.LazyFrame:
        """run all steps to convert collated data to tokenized timelines"""
        df = self.get_data()  # load data
        df = self.add_ends(df)  # add BOS/EOS tokens
        df = self.add_clocks(df)  # add clock tokens if configured
        df = self.bin_data(df)  # create bins from training data and bin numeric values
        df = self.insert_time_spacers(df)  # insert time spacer codes when configured
        df = self.get_pretokenized(df).cache()  # pretokenize
        df = self.tokenize_data(df)  # collect tokens into timelines

        if verbose:
            self.logger.summarize_tokens_times(df, self.subject_splits, self.lookup)

        return df

    def save_all(self, verbose: bool = False):
        """
        get tokenized timelines and save them to disc,
        along with artifacts created during tokenization
        """
        df = self.get_all(verbose)
        df.sink_parquet(
            self.processed_data_home / "tokens_times.parquet", engine="streaming"
        )
        self.save(self.processed_data_home / "tokenizer.yaml")

    def __call__(self, word: str) -> int:
        """apply tokenizer to a single word"""
        try:
            return (
                self.lookup.filter(pl.col("to_tokenize") == word).select("token").item()
            )
        except (ValueError, AttributeError):
            return 0  # UNK

    def __contains__(self, word: str) -> bool:
        """is word in the tokenizer's vocabulary?"""
        return self.__call__(word) != 0

    def __str__(self):
        return "{sp} of {sz} words {md}".format(
            sp=self.__class__,
            sz=len(self),
            md="in training mode" if self.is_training else "(frozen)",
        )

    def __repr__(self):
        return str(self) + ", created {dttm}".format(dttm=self.created_dttm)

    def __len__(self) -> int:
        """number of tokens in vocabulary"""
        return len(self.lookup) if self.lookup is not None else 0

    def to_yaml(self) -> str:
        """yaml representation of tokenizer; sufficient for reconstruction"""
        return OmegaConf.to_yaml(
            {
                "lookup": dict(self.lookup.rows()) if self.lookup is not None else None,
                "bins": {k: v for k, *v in self.bins.rows()}
                if self.bins is not None
                else None,
                "is_training": self.is_training,
                "cfg": OmegaConf.to_container(self.cfg),
                "created_dttm": self.created_dttm,
                "cocoa_version": meta.version("cocoa"),
            }
        )

    def from_yaml(self, yaml_str: str, done_training=True) -> "Tokenizer":
        """
        construct tokenizer from yaml representation
        places tokenizer into inference mode by default
        """
        data = OmegaConf.create(yaml_str)
        cfg = OmegaConf.to_container(data.cfg)
        tkzr = self.__class__(
            self.config_file,
            processed_data_home=self.processed_data_home,
            is_training=data.is_training,
            **OmegaConf.merge(self.cfg, cfg),
        )
        tkzr.created_dttm = data.created_dttm
        if data.bins is not None:
            tkzr.bins = pl.DataFrame(
                [[k, *v] for k, v in dict(data.bins).items()],
                schema=["code", *[f"break_{i}" for i in range(1, cfg["n_bins"])]],
                orient="row",
            )
        if data.lookup is not None:
            tkzr.lookup = pl.DataFrame(
                list(dict(data.lookup).items()),
                schema={"to_tokenize": pl.String, "token": pl.UInt32},
                orient="row",
            )
        if done_training:
            tkzr.is_training = False
        return tkzr

    def save(self, path: pathlib.Path):
        """write yaml representation of tokenizer to disc"""
        to_file = pathlib.Path(path).expanduser().resolve()
        to_file.parent.mkdir(parents=True, exist_ok=True)
        with open(to_file, "w") as f:
            f.write(self.to_yaml())

    def load(self, path: pathlib.Path, done_training=True):
        """retrieve tokenizer from saved yaml representation"""
        from_file = pathlib.Path(path).expanduser().resolve()
        with open(from_file, "r") as f:
            yaml_str = f.read()
        return self.from_yaml(yaml_str, done_training=done_training)

__call__(word)

apply tokenizer to a single word

Source code in src/cocoa/tokenizer.py

def __call__(self, word: str) -> int:
    """apply tokenizer to a single word"""
    try:
        return (
            self.lookup.filter(pl.col("to_tokenize") == word).select("token").item()
        )
    except (ValueError, AttributeError):
        return 0  # UNK

__contains__(word)

is word in the tokenizer's vocabulary?

Source code in src/cocoa/tokenizer.py

def __contains__(self, word: str) -> bool:
    """is word in the tokenizer's vocabulary?"""
    return self.__call__(word) != 0

__len__()

number of tokens in vocabulary

Source code in src/cocoa/tokenizer.py

def __len__(self) -> int:
    """number of tokens in vocabulary"""
    return len(self.lookup) if self.lookup is not None else 0

add_clocks(df)

add clock codes if configured

Source code in src/cocoa/tokenizer.py

def add_clocks(self, df: pl.LazyFrame) -> pl.LazyFrame:
    """add clock codes if configured"""
    if self.cfg.get("insert_clocks", False):
        return pl.concat(
            [
                df,
                df.group_by("subject_id")
                .agg(
                    start=pl.col("time").min().dt.truncate("1h"),
                    end=pl.col("time").max().dt.truncate("1h"),
                )
                .with_columns(
                    pl.datetime_ranges(
                        pl.col("start"),
                        pl.col("end"),
                        interval="1h",
                        closed="right",
                    )
                    .alias("time")
                    .list.filter(
                        pl.element().dt.strftime("%H").is_in(list(self.cfg.clocks))
                    )
                )
                .explode("time", keep_nulls=False)
                .drop_nulls(subset=["time"])
                .with_columns(HH=pl.col("time").dt.strftime("%H"))
                .select(
                    "subject_id",
                    "time",
                    pl.concat_str(pl.lit("CLCK//"), pl.col("HH")).alias("code"),
                ),
            ],
            how="diagonal",
        )
    else:
        return df

add_ends(df) staticmethod

add BOS / EOS codes at appropriate places

Source code in src/cocoa/tokenizer.py

@staticmethod
def add_ends(df: pl.LazyFrame) -> pl.LazyFrame:
    """add BOS / EOS codes at appropriate places"""
    return pl.concat(
        [
            df,
            df.group_by("subject_id")
            .agg(pl.col("time").min())
            .with_columns(code=pl.lit("BOS")),
            df.group_by("subject_id")
            .agg(pl.col("time").max())
            .with_columns(code=pl.lit("EOS")),
        ],
        how="diagonal",
    )

bin_data(df)

discretize numeric values with learned cut points

Source code in src/cocoa/tokenizer.py

def bin_data(self, df: pl.LazyFrame) -> pl.LazyFrame:
    """discretize numeric values with learned cut points"""
    return (
        df.join(self.get_bins(df).lazy(), on="code", how="left")
        .with_columns(
            pl.when(pl.col("numeric_value").is_not_null())
            .then(
                pl.concat_str(
                    pl.lit("Q"),
                    pl.sum_horizontal(
                        [
                            pl.col(f"break_{i}") <= pl.col("numeric_value")
                            for i in range(1, self.cfg.n_bins)
                        ]
                    ).cast(pl.String),
                )
            )
            .otherwise(None)
            .alias("binned_value")
        )
        .drop([f"break_{i}" for i in range(1, self.cfg.n_bins)])
    )

from_yaml(yaml_str, done_training=True)

construct tokenizer from yaml representation places tokenizer into inference mode by default

Source code in src/cocoa/tokenizer.py

def from_yaml(self, yaml_str: str, done_training=True) -> "Tokenizer":
    """
    construct tokenizer from yaml representation
    places tokenizer into inference mode by default
    """
    data = OmegaConf.create(yaml_str)
    cfg = OmegaConf.to_container(data.cfg)
    tkzr = self.__class__(
        self.config_file,
        processed_data_home=self.processed_data_home,
        is_training=data.is_training,
        **OmegaConf.merge(self.cfg, cfg),
    )
    tkzr.created_dttm = data.created_dttm
    if data.bins is not None:
        tkzr.bins = pl.DataFrame(
            [[k, *v] for k, v in dict(data.bins).items()],
            schema=["code", *[f"break_{i}" for i in range(1, cfg["n_bins"])]],
            orient="row",
        )
    if data.lookup is not None:
        tkzr.lookup = pl.DataFrame(
            list(dict(data.lookup).items()),
            schema={"to_tokenize": pl.String, "token": pl.UInt32},
            orient="row",
        )
    if done_training:
        tkzr.is_training = False
    return tkzr

get_all(verbose=False)

run all steps to convert collated data to tokenized timelines

Source code in src/cocoa/tokenizer.py

def get_all(self, verbose: bool = False) -> pl.LazyFrame:
    """run all steps to convert collated data to tokenized timelines"""
    df = self.get_data()  # load data
    df = self.add_ends(df)  # add BOS/EOS tokens
    df = self.add_clocks(df)  # add clock tokens if configured
    df = self.bin_data(df)  # create bins from training data and bin numeric values
    df = self.insert_time_spacers(df)  # insert time spacer codes when configured
    df = self.get_pretokenized(df).cache()  # pretokenize
    df = self.tokenize_data(df)  # collect tokens into timelines

    if verbose:
        self.logger.summarize_tokens_times(df, self.subject_splits, self.lookup)

    return df

get_bins(df)

calculate bins for numeric values on training data

Source code in src/cocoa/tokenizer.py

def get_bins(self, df: pl.LazyFrame) -> pl.DataFrame:
    """calculate bins for numeric values on training data"""
    if self.bins is None:
        assert self.is_training, "Bins must be learned during training"
        self.bins = (
            df.join(  # restrict to training data
                self.subject_splits.filter(pl.col("split") == "train"),
                on="subject_id",
                validate="m:1",
            )
            .drop_nulls(subset=["numeric_value"])
            .drop_nans(subset=["numeric_value"])
            .group_by("code")
            .agg(
                [
                    pl.col("numeric_value")
                    .quantile(i / self.cfg.n_bins)
                    .alias(f"break_{i}")
                    for i in range(1, self.cfg.n_bins)
                ]
            )
        ).collect()
    return self.bins

get_lookup(pt)

create mapping from vocabulary to integer tokens on training data

Source code in src/cocoa/tokenizer.py

def get_lookup(self, pt: pl.LazyFrame) -> pl.DataFrame:
    """create mapping from vocabulary to integer tokens on training data"""
    if self.lookup is None:
        assert self.is_training, "Lookup table must be learned during training"
        lookup = (
            pt.join(  # restrict to training data
                self.subject_splits.filter(pl.col("split") == "train"),
                on="subject_id",
                validate="m:1",
            )
            .explode("to_tokenize")
            .select(pl.col("to_tokenize").unique().sort())
            .filter(pl.col("to_tokenize") != "UNK")  # UNK is 0
            .with_row_index("token", offset=1)
            .select("to_tokenize", "token")
        )
        unk_row = pl.LazyFrame(
            {"to_tokenize": ["UNK"], "token": pl.Series([0], dtype=pl.UInt32)}
        )
        self.lookup = pl.concat([unk_row, lookup]).collect()
    return self.lookup

get_pretokenized(df)

prepare codes for tokenization, depending on whether fusion is configured

Source code in src/cocoa/tokenizer.py

def get_pretokenized(self, df: pl.LazyFrame) -> pl.LazyFrame:
    """prepare codes for tokenization, depending on whether fusion is configured"""
    return df.with_columns(
        pl.concat_list(
            pl.col("t_spacer"),
            pl.concat_str(
                pl.col("code"),
                pl.col("binned_value"),
                pl.col("text_value"),
                separator="_",
                ignore_nulls=True,
            )
            if self.cfg.get("fused", False)
            else pl.concat_list("code", "binned_value", "text_value"),
        )
        .list.drop_nulls()
        .alias("to_tokenize")
    )

get_priority()

fetch ordering for sorting cotemporaneous codes

Source code in src/cocoa/tokenizer.py

def get_priority(self) -> pl.DataFrame:
    """fetch ordering for sorting cotemporaneous codes"""
    return (
        pl.Series("code_type", self.cfg.ordering)
        .to_frame()
        .with_row_index("priority")
    )

insert_time_spacers(df)

add time spacing tokens if configured; this should be done after clock tokens are inserted

Source code in src/cocoa/tokenizer.py

def insert_time_spacers(self, df: pl.LazyFrame) -> pl.LazyFrame:
    """
    add time spacing tokens if configured;
    this should be done after clock tokens are inserted
    """
    if self.cfg.get("insert_spacers", False):
        spcrs = dict(self.cfg.spacers)
        return (
            df.sort("subject_id", "time")
            .with_columns(
                tdiff_mins=(
                    pl.col("time") - pl.col("time").shift(1).over("subject_id")
                ).dt.total_minutes()
            )
            .with_columns(
                t_spacer=pl.when(
                    (pl.col("tdiff_mins") < min(spcrs.values()))
                    | pl.col("tdiff_mins").is_null()
                )
                .then(None)
                .otherwise(
                    pl.concat_str(
                        pl.lit("TIME//"),
                        pl.col("tdiff_mins")
                        .cut(
                            breaks=list(spcrs.values())[1:],
                            labels=list(spcrs.keys()),
                        )
                        .cast(pl.String),
                    )
                )
                .cast(pl.String)
            )
        )
    else:
        return df.with_columns(t_spacer=pl.lit(None))

load(path, done_training=True)

retrieve tokenizer from saved yaml representation

Source code in src/cocoa/tokenizer.py

def load(self, path: pathlib.Path, done_training=True):
    """retrieve tokenizer from saved yaml representation"""
    from_file = pathlib.Path(path).expanduser().resolve()
    with open(from_file, "r") as f:
        yaml_str = f.read()
    return self.from_yaml(yaml_str, done_training=done_training)

save(path)

write yaml representation of tokenizer to disc

Source code in src/cocoa/tokenizer.py

def save(self, path: pathlib.Path):
    """write yaml representation of tokenizer to disc"""
    to_file = pathlib.Path(path).expanduser().resolve()
    to_file.parent.mkdir(parents=True, exist_ok=True)
    with open(to_file, "w") as f:
        f.write(self.to_yaml())

save_all(verbose=False)

get tokenized timelines and save them to disc, along with artifacts created during tokenization

Source code in src/cocoa/tokenizer.py

def save_all(self, verbose: bool = False):
    """
    get tokenized timelines and save them to disc,
    along with artifacts created during tokenization
    """
    df = self.get_all(verbose)
    df.sink_parquet(
        self.processed_data_home / "tokens_times.parquet", engine="streaming"
    )
    self.save(self.processed_data_home / "tokenizer.yaml")

to_yaml()

yaml representation of tokenizer; sufficient for reconstruction

Source code in src/cocoa/tokenizer.py

def to_yaml(self) -> str:
    """yaml representation of tokenizer; sufficient for reconstruction"""
    return OmegaConf.to_yaml(
        {
            "lookup": dict(self.lookup.rows()) if self.lookup is not None else None,
            "bins": {k: v for k, *v in self.bins.rows()}
            if self.bins is not None
            else None,
            "is_training": self.is_training,
            "cfg": OmegaConf.to_container(self.cfg),
            "created_dttm": self.created_dttm,
            "cocoa_version": meta.version("cocoa"),
        }
    )

tokenize_data(pt)

apply lookup table to pretokenized data

Source code in src/cocoa/tokenizer.py

def tokenize_data(self, pt: pl.LazyFrame) -> pl.LazyFrame:
    """apply lookup table to pretokenized data"""
    return (
        pt.with_columns(code_type=pl.col("code").str.split("//").list[0])
        .join(
            self.get_priority().lazy(), on="code_type", how="left", validate="m:1"
        )
        .with_columns(pl.col("priority").fill_null(len(self.cfg.ordering)))
        .sort("time", "priority")
        .explode("to_tokenize")
        .join(
            self.get_lookup(pt).lazy(), on="to_tokenize", validate="m:1", how="left"
        )
        .with_columns(
            pl.col("token").fill_null(pl.lit(0, dtype=pl.UInt32))
        )  # UNK is 0
        .group_by("subject_id", maintain_order=True)
        .agg(
            pl.col("token").alias("tokens"),
            pl.col("time").alias("times"),
            *(
                [pl.col("numeric_value").alias("numeric_values")]
                if self.cfg.get("include_numeric_values", False)
                else []
            ),
        )
    )