Skip to content

custom_text_tensorization

Functions for tensorizing MEDS datasets.

convert_to_NRT(df)

This converts a tokenized dataframe into a nested ragged tensor.

Most of the work for this function is actually done in tokenize – this function is just a wrapper to convert the output into a nested ragged tensor using polars’ built-in to_dict method.

Parameters:

Name Type Description Default
df LazyFrame

The tokenized dataframe.

 required

Returns:

Type Description
JointNestedRaggedTensorDict

A JointNestedRaggedTensorDict object representing the tokenized dataframe, accounting for however
JointNestedRaggedTensorDict

many levels of ragged nesting are present among the codes and numeric values.

Raises:

Type Description
ValueError

If there are no time delta columns or if there are multiple time delta columns.

Examples:

>>> df = pl.DataFrame({
...     "subject_id": [1, 2],
...     "time_delta_days": [[float("nan"), 12.0], [float("nan")]],
...     "code": [[[101, 102], [103]], [[201, 202]]],
...     "modality_idx": [[[0., float("nan")], [2.]], [[float("nan"), float("nan")]]],
...     "numeric_value": [[[2.0, 3.0], [4.0]], [[6.0, 7.0]]]
... })
>>> nrt = convert_to_NRT(df.lazy())
>>> for k, v in sorted(list(nrt.to_dense().items())):
...     print(k)
...     print(v)
code
[[[101 102]
  [103   0]]

 [[201 202]
  [  0   0]]]
dim1/mask
[[ True  True]
 [ True False]]
dim2/mask
[[[ True  True]
  [ True False]]

 [[ True  True]
  [False False]]]
modality_idx
[[[ 0. nan]
  [ 2.  0.]]

 [[nan nan]
  [ 0.  0.]]]
numeric_value
[[[2. 3.]
  [4. 0.]]

 [[6. 7.]
  [0. 0.]]]
time_delta_days
[[nan 12.]
 [nan  0.]]
Source code in meds_torch/utils/custom_text_tensorization.py 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def convert_to_NRT(df: pl.LazyFrame) -> JointNestedRaggedTensorDict:
    """This converts a tokenized dataframe into a nested ragged tensor.

    Most of the work for this function is actually done in `tokenize` -- this function is just a wrapper
    to convert the output into a nested ragged tensor using polars' built-in `to_dict` method.

    Args:
        df: The tokenized dataframe.

    Returns:
        A `JointNestedRaggedTensorDict` object representing the tokenized dataframe, accounting for however
        many levels of ragged nesting are present among the codes and numeric values.

    Raises:
        ValueError: If there are no time delta columns or if there are multiple time delta columns.

    Examples:
        >>> df = pl.DataFrame({
        ...     "subject_id": [1, 2],
        ...     "time_delta_days": [[float("nan"), 12.0], [float("nan")]],
        ...     "code": [[[101, 102], [103]], [[201, 202]]],
        ...     "modality_idx": [[[0., float("nan")], [2.]], [[float("nan"), float("nan")]]],
        ...     "numeric_value": [[[2.0, 3.0], [4.0]], [[6.0, 7.0]]]
        ... })
        >>> nrt = convert_to_NRT(df.lazy())
        >>> for k, v in sorted(list(nrt.to_dense().items())):
        ...     print(k)
        ...     print(v)
        code
        [[[101 102]
          [103   0]]
        <BLANKLINE>
         [[201 202]
          [  0   0]]]
        dim1/mask
        [[ True  True]
         [ True False]]
        dim2/mask
        [[[ True  True]
          [ True False]]
        <BLANKLINE>
         [[ True  True]
          [False False]]]
        modality_idx
        [[[ 0. nan]
          [ 2.  0.]]
        <BLANKLINE>
         [[nan nan]
          [ 0.  0.]]]
        numeric_value
        [[[2. 3.]
          [4. 0.]]
        <BLANKLINE>
         [[6. 7.]
          [0. 0.]]]
        time_delta_days
        [[nan 12.]
         [nan  0.]]
    """

    # There should only be one time delta column, but this ensures we catch it regardless of the unit of time
    # used to convert the time deltas, and that we verify there is only one such column.
    time_delta_cols = [c for c in df.collect_schema().names() if c.startswith("time_delta_")]

    if len(time_delta_cols) == 0:
        raise ValueError("Expected at least one time delta column, found none")
    elif len(time_delta_cols) > 1:
        raise ValueError(f"Expected exactly one time delta column, found columns: {time_delta_cols}")

    time_delta_col = time_delta_cols[0]

    tensors_dict = (
        df.select(time_delta_col, "code", "numeric_value", pl.col("modality_idx"))
        .collect()
        .to_dict(as_series=False)
    )

    if all((not v) for v in tensors_dict.values()):
        logger.warning("All columns are empty. Returning an empty tensor dict.")
        return JointNestedRaggedTensorDict({})

    for k, v in tensors_dict.items():
        if not v:
            raise ValueError(f"Column {k} is empty")

    return JointNestedRaggedTensorDict(tensors_dict)

main(cfg)

TODO.

Source code in meds_torch/utils/custom_text_tensorization.py 
104
105
106
107
108
109
110
111
112
113
114
115
@hydra.main(
    version_base=None, config_path=str(PREPROCESS_CONFIG_YAML.parent), config_name=PREPROCESS_CONFIG_YAML.stem
)
def main(cfg: DictConfig):
    """TODO."""

    map_over(
        cfg,
        compute_fn=convert_to_NRT,
        write_fn=JointNestedRaggedTensorDict.save,
        shard_iterator_fntr=partial(shard_iterator, in_prefix="event_seqs/", out_suffix=".nrt"),
    )