From de9227badd8b042beadd4bfa118fe75531fe7f59 Mon Sep 17 00:00:00 2001 From: Germans Savcisens <38817889+carlomarxdk@users.noreply.github.com> Date: Wed, 17 Jun 2026 14:46:22 -0400 Subject: [PATCH] Support for the full static sources --- docs/getting-started/quickstart.md | 52 +++ docs/guide/sources.md | 32 +- examples/synthetic_source.ipynb | 563 ++++++++++++++++++++++++++--- pyproject.toml | 2 +- src/tab2seq/source/config.py | 18 +- 5 files changed, 606 insertions(+), 61 deletions(-) diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md index a154595..053cf2a 100644 --- a/docs/getting-started/quickstart.md +++ b/docs/getting-started/quickstart.md @@ -33,6 +33,58 @@ shape: (5, 7) Each `Source` describes one event table: its file path, ID column, timestamp, and feature columns. +### Define one by one + +```python +## Option 1: Source(SourceConfig) +## Example with the health data +source_H = Source(config=SourceConfig( + name="health", + filepath="synthetic_data/health.parquet", + id_col="entity_id", + categorical_cols=[ + CategoricalColConfig(col_name="diagnosis", prefix="DIAG"), + CategoricalColConfig(col_name="procedure", prefix="PROC"), + CategoricalColConfig(col_name="department", prefix="DEPT"), + ], + continuous_cols=[ + ContinuousColConfig(col_name="cost", prefix="COST", n_bins=20, strategy="quantile"), + ContinuousColConfig(col_name="length_of_stay", prefix="LOS", n_bins=20, strategy="quantile"), + ], + temporal_cols=[ + TemporalColConfig(col_name="date", is_primary=True, drop_na=True, col_type="datetime") + ], + output_format="parquet", +)) + +## Option 2: SourceConfig -> Source +## Example with the labor data +config_L = SourceConfig( + name="labour", + filepath="synthetic_data/labour.parquet", + id_col="entity_id", + categorical_cols=[ + CategoricalColConfig(col_name="status", prefix="STATUS"), + CategoricalColConfig(col_name="occupation", prefix="OCC"), + CategoricalColConfig(col_name="residence_region", prefix="REGION"), + CategoricalColConfig(col_name="native_language", prefix="LANG", static=True), + ], + continuous_cols=[ + ContinuousColConfig(col_name="weekly_hours", prefix="WEEKLY_HOURS") + ], + temporal_cols=[ + TemporalColConfig(col_name="date", is_primary=True, drop_na=True, col_type="datetime"), + TemporalColConfig(col_name="birthday", is_primary=False, static=True, drop_na=True, col_type="datetime"), + ], + output_format="parquet", +) +source_L = Source(config=config_L) +``` + +You can then pass the sources to the `Cohort` object as a list `[source_H, source_L]`. + +### Define via `Source Collection` + ```python from tab2seq.source import ( Source, SourceCollection, SourceConfig, diff --git a/docs/guide/sources.md b/docs/guide/sources.md index 980b985..18684b8 100644 --- a/docs/guide/sources.md +++ b/docs/guide/sources.md @@ -18,7 +18,33 @@ Columns marked `static=True` represent entity-level attributes that do not chang - Carried through to the cohort split table as entity attributes - Available as input to `RelativeDateRule` for computing relative-date features -## Defining a SourceCollection +## Defining a Single Source + +```python +## example with the Labor data +config_L = SourceConfig( + name="labour", + filepath="synthetic_data/labour.parquet", + id_col="entity_id", + categorical_cols=[ + CategoricalColConfig(col_name="status", prefix="STATUS"), + CategoricalColConfig(col_name="occupation", prefix="OCC"), + CategoricalColConfig(col_name="residence_region", prefix="REGION"), + CategoricalColConfig(col_name="native_language", prefix="LANG", static=True), + ], + continuous_cols=[ + ContinuousColConfig(col_name="weekly_hours", prefix="WEEKLY_HOURS") + ], + temporal_cols=[ + TemporalColConfig(col_name="date", is_primary=True, drop_na=True, col_type="datetime"), + TemporalColConfig(col_name="birthday", is_primary=False, static=True, drop_na=True, col_type="datetime"), + ], + output_format="parquet", +) +source_L = Source(config=config_L) +``` + +## Defining via SourceCollection ```python from tab2seq.source import ( @@ -58,3 +84,7 @@ Bin edges are fitted on train data only and serialised with the vocabulary. ## Prefixes Each column config has a `prefix` that becomes the token string prefix, e.g. `DIAG_J18.1`, `COST_bin_3`. Prefixes must be unique within a source. + +## Source with (Only Static Attributes) + +If you have a source with only static attributes (aka no `temporal_cols`), you still can uase this object. Do not forget to assign `static=True` to every feature (otherwise you will get an error). diff --git a/examples/synthetic_source.ipynb b/examples/synthetic_source.ipynb index 96415e9..017c202 100644 --- a/examples/synthetic_source.ipynb +++ b/examples/synthetic_source.ipynb @@ -31,7 +31,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "5970bb01", "metadata": {}, "outputs": [], @@ -42,10 +42,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "e5fadc97", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Generated synthetic data at: {'health': PosixPath('synthetic_data/health.parquet'), 'labour': PosixPath('synthetic_data/labour.parquet'), 'survey': PosixPath('synthetic_data/survey.parquet'), 'income': PosixPath('synthetic_data/income.parquet')}\n" + ] + } + ], "source": [ "data_paths = generate_synthetic_data(output_dir=\"synthetic_data\", \n", " n_entities=10000, \n", @@ -80,10 +88,42 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "ae32a281", "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "shape: (5, 7)
entity_iddatediagnosisproceduredepartmentcostlength_of_stay
strdatestrstrstrf64i64
"E00001"2016-09-15"J18.1""CABG""gastroenterology"7306.172
"E00001"2017-05-25"E78.0""XRAY""neurology"138.651
"E00001"2018-01-18"E78.0""MRI""general_surgery"6704.5910
"E00001"2019-11-11"C34.1""ECHO""general_surgery"910.120
"E00001"2020-05-20"E78.0""DIALYSIS""neurology"2266.522
" + ], + "text/plain": [ + "shape: (5, 7)\n", + "┌───────────┬────────────┬───────────┬───────────┬──────────────────┬─────────┬────────────────┐\n", + "│ entity_id ┆ date ┆ diagnosis ┆ procedure ┆ department ┆ cost ┆ length_of_stay │\n", + "│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n", + "│ str ┆ date ┆ str ┆ str ┆ str ┆ f64 ┆ i64 │\n", + "╞═══════════╪════════════╪═══════════╪═══════════╪══════════════════╪═════════╪════════════════╡\n", + "│ E00001 ┆ 2016-09-15 ┆ J18.1 ┆ CABG ┆ gastroenterology ┆ 7306.17 ┆ 2 │\n", + "│ E00001 ┆ 2017-05-25 ┆ E78.0 ┆ XRAY ┆ neurology ┆ 138.65 ┆ 1 │\n", + "│ E00001 ┆ 2018-01-18 ┆ E78.0 ┆ MRI ┆ general_surgery ┆ 6704.59 ┆ 10 │\n", + "│ E00001 ┆ 2019-11-11 ┆ C34.1 ┆ ECHO ┆ general_surgery ┆ 910.12 ┆ 0 │\n", + "│ E00001 ┆ 2020-05-20 ┆ E78.0 ┆ DIALYSIS ┆ neurology ┆ 2266.52 ┆ 2 │\n", + "└───────────┴────────────┴───────────┴───────────┴──────────────────┴─────────┴────────────────┘" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "lf_health = pl.read_parquet(data_paths[\"health\"])\n", "lf_health.head()" @@ -91,10 +131,58 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "e8d70d21", "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "shape: (10, 8)
entity_iddatestatusoccupationweekly_hoursresidence_regionbirthdaynative_language
strdatestrstrf64strdatestr
"E07897"2018-11-01"student""agriculture"0.0"central"1951-04-20"spanish"
"E08737"2019-02-01"self_employed""manufacturing"35.9"capital"1972-10-11"swahili"
"E08992"2022-08-01"self_employed""engineering"31.9"east"1971-10-03"arabic"
"E03000"2023-04-01"self_employed""transport"38.9"east"1954-03-07"turkish"
"E05609"2023-07-01"student""IT"0.0"central"1982-08-12"portuguese"
"E03553"2016-07-01"parental_leave""construction"0.0"west"1975-05-09"swahili"
"E02355"2015-08-01"employed""construction"43.2"capital"1964-02-20"hindi"
"E02435"2024-01-01"student""IT"0.0"island"1986-02-15"french"
"E05554"2015-02-01"student""agriculture"0.0"capital"1954-07-18"arabic"
"E06737"2021-02-01"unemployed""education"0.0"north"1968-07-14"swahili"
" + ], + "text/plain": [ + "shape: (10, 8)\n", + "┌───────────┬────────────┬────────────┬────────────┬───────────┬───────────┬───────────┬───────────┐\n", + "│ entity_id ┆ date ┆ status ┆ occupation ┆ weekly_ho ┆ residence ┆ birthday ┆ native_la │\n", + "│ --- ┆ --- ┆ --- ┆ --- ┆ urs ┆ _region ┆ --- ┆ nguage │\n", + "│ str ┆ date ┆ str ┆ str ┆ --- ┆ --- ┆ date ┆ --- │\n", + "│ ┆ ┆ ┆ ┆ f64 ┆ str ┆ ┆ str │\n", + "╞═══════════╪════════════╪════════════╪════════════╪═══════════╪═══════════╪═══════════╪═══════════╡\n", + "│ E07897 ┆ 2018-11-01 ┆ student ┆ agricultur ┆ 0.0 ┆ central ┆ 1951-04-2 ┆ spanish │\n", + "│ ┆ ┆ ┆ e ┆ ┆ ┆ 0 ┆ │\n", + "│ E08737 ┆ 2019-02-01 ┆ self_emplo ┆ manufactur ┆ 35.9 ┆ capital ┆ 1972-10-1 ┆ swahili │\n", + "│ ┆ ┆ yed ┆ ing ┆ ┆ ┆ 1 ┆ │\n", + "│ E08992 ┆ 2022-08-01 ┆ self_emplo ┆ engineerin ┆ 31.9 ┆ east ┆ 1971-10-0 ┆ arabic │\n", + "│ ┆ ┆ yed ┆ g ┆ ┆ ┆ 3 ┆ │\n", + "│ E03000 ┆ 2023-04-01 ┆ self_emplo ┆ transport ┆ 38.9 ┆ east ┆ 1954-03-0 ┆ turkish │\n", + "│ ┆ ┆ yed ┆ ┆ ┆ ┆ 7 ┆ │\n", + "│ E05609 ┆ 2023-07-01 ┆ student ┆ IT ┆ 0.0 ┆ central ┆ 1982-08-1 ┆ portugues │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 2 ┆ e │\n", + "│ E03553 ┆ 2016-07-01 ┆ parental_l ┆ constructi ┆ 0.0 ┆ west ┆ 1975-05-0 ┆ swahili │\n", + "│ ┆ ┆ eave ┆ on ┆ ┆ ┆ 9 ┆ │\n", + "│ E02355 ┆ 2015-08-01 ┆ employed ┆ constructi ┆ 43.2 ┆ capital ┆ 1964-02-2 ┆ hindi │\n", + "│ ┆ ┆ ┆ on ┆ ┆ ┆ 0 ┆ │\n", + "│ E02435 ┆ 2024-01-01 ┆ student ┆ IT ┆ 0.0 ┆ island ┆ 1986-02-1 ┆ french │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 5 ┆ │\n", + "│ E05554 ┆ 2015-02-01 ┆ student ┆ agricultur ┆ 0.0 ┆ capital ┆ 1954-07-1 ┆ arabic │\n", + "│ ┆ ┆ ┆ e ┆ ┆ ┆ 8 ┆ │\n", + "│ E06737 ┆ 2021-02-01 ┆ unemployed ┆ education ┆ 0.0 ┆ north ┆ 1968-07-1 ┆ swahili │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 4 ┆ │\n", + "└───────────┴────────────┴────────────┴────────────┴───────────┴───────────┴───────────┴───────────┘" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "lf_labour = pl.read_parquet(data_paths[\"labour\"])\n", "lf_labour.sample(10)" @@ -125,7 +213,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "ea2afe64", "metadata": {}, "outputs": [], @@ -142,10 +230,49 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 6, "id": "6e9fc417", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Unique entities in health source: 9797\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "shape: (5, 7)
entity_iddatediagnosisproceduredepartmentcostlength_of_stay
strdatestrstrstrf64i64
"E00001"2016-09-15"J18.1""CABG""gastroenterology"7306.172
"E00001"2017-05-25"E78.0""XRAY""neurology"138.651
"E00001"2018-01-18"E78.0""MRI""general_surgery"6704.5910
"E00001"2019-11-11"C34.1""ECHO""general_surgery"910.120
"E00001"2020-05-20"E78.0""DIALYSIS""neurology"2266.522
" + ], + "text/plain": [ + "shape: (5, 7)\n", + "┌───────────┬────────────┬───────────┬───────────┬──────────────────┬─────────┬────────────────┐\n", + "│ entity_id ┆ date ┆ diagnosis ┆ procedure ┆ department ┆ cost ┆ length_of_stay │\n", + "│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n", + "│ str ┆ date ┆ str ┆ str ┆ str ┆ f64 ┆ i64 │\n", + "╞═══════════╪════════════╪═══════════╪═══════════╪══════════════════╪═════════╪════════════════╡\n", + "│ E00001 ┆ 2016-09-15 ┆ J18.1 ┆ CABG ┆ gastroenterology ┆ 7306.17 ┆ 2 │\n", + "│ E00001 ┆ 2017-05-25 ┆ E78.0 ┆ XRAY ┆ neurology ┆ 138.65 ┆ 1 │\n", + "│ E00001 ┆ 2018-01-18 ┆ E78.0 ┆ MRI ┆ general_surgery ┆ 6704.59 ┆ 10 │\n", + "│ E00001 ┆ 2019-11-11 ┆ C34.1 ┆ ECHO ┆ general_surgery ┆ 910.12 ┆ 0 │\n", + "│ E00001 ┆ 2020-05-20 ┆ E78.0 ┆ DIALYSIS ┆ neurology ┆ 2266.52 ┆ 2 │\n", + "└───────────┴────────────┴───────────┴───────────┴──────────────────┴─────────┴────────────────┘" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "source_H = Source(config=SourceConfig(\n", " name=\"health\",\n", @@ -174,10 +301,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "id": "bb8c553d", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Number of unique IDs: 10000\n" + ] + } + ], "source": [ "# or you could define the Source config separately and then create the Source\n", "\n", @@ -207,10 +342,73 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 8, + "id": "e87d80ca", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "shape: (5, 8)
entity_iddatebirthdaystatusoccupationresidence_regionnative_languageweekly_hours
strdatedatestrstrstrstrf64
"E00001"2015-01-011954-11-11"employed""education""capital""spanish"36.4
"E00001"2015-03-011954-11-11"employed""education""capital""hindi"36.9
"E00001"2015-05-011954-11-11"employed""education""capital""german"30.7
"E00001"2015-08-011954-11-11"employed""education""capital""german"34.2
"E00001"2015-11-011954-11-11"employed""education""capital""portuguese"38.5
" + ], + "text/plain": [ + "shape: (5, 8)\n", + "┌───────────┬────────────┬────────────┬──────────┬────────────┬────────────┬───────────┬───────────┐\n", + "│ entity_id ┆ date ┆ birthday ┆ status ┆ occupation ┆ residence_ ┆ native_la ┆ weekly_ho │\n", + "│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ region ┆ nguage ┆ urs │\n", + "│ str ┆ date ┆ date ┆ str ┆ str ┆ --- ┆ --- ┆ --- │\n", + "│ ┆ ┆ ┆ ┆ ┆ str ┆ str ┆ f64 │\n", + "╞═══════════╪════════════╪════════════╪══════════╪════════════╪════════════╪═══════════╪═══════════╡\n", + "│ E00001 ┆ 2015-01-01 ┆ 1954-11-11 ┆ employed ┆ education ┆ capital ┆ spanish ┆ 36.4 │\n", + "│ E00001 ┆ 2015-03-01 ┆ 1954-11-11 ┆ employed ┆ education ┆ capital ┆ hindi ┆ 36.9 │\n", + "│ E00001 ┆ 2015-05-01 ┆ 1954-11-11 ┆ employed ┆ education ┆ capital ┆ german ┆ 30.7 │\n", + "│ E00001 ┆ 2015-08-01 ┆ 1954-11-11 ┆ employed ┆ education ┆ capital ┆ german ┆ 34.2 │\n", + "│ E00001 ┆ 2015-11-01 ┆ 1954-11-11 ┆ employed ┆ education ┆ capital ┆ portugues ┆ 38.5 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ e ┆ │\n", + "└───────────┴────────────┴────────────┴──────────┴────────────┴────────────┴───────────┴───────────┘" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "source_L.read_all().head()" + ] + }, + { + "cell_type": "code", + "execution_count": 9, "id": "2d0ae9bd", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "All unique entity IDs in collection: 10000\n" + ] + }, + { + "data": { + "text/plain": [ + "{'health': Source(name='health', path=PosixPath('synthetic_data/health.parquet')),\n", + " 'labour': Source(name='labour', path=PosixPath('synthetic_data/labour.parquet'))}" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# You can also create a SourceCollection to manage multiple sources together\n", "collection = SourceCollection(sources=[source_H, source_L])\n", @@ -223,10 +421,19 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 10, "id": "40d6afc7", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "health: 9797 entities\n", + "labour: 10000 entities\n" + ] + } + ], "source": [ "# Or you can make collections directly from configs\n", "from tab2seq.source import (\n", @@ -323,10 +530,50 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 11, "id": "644bc2a2", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Entities in cohort: 10000\n", + "The train/val/test split data also keeps track of the static features\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "shape: (5, 4)
entity_idlabour__birthdaylabour__native_languagesplit
strdatestrstr
"E00001"1954-11-11"spanish""val"
"E00002"2004-04-16"japanese""val"
"E00003"1950-08-17"english""train"
"E00004"1992-06-23"russian""val"
"E00005"1964-04-15"japanese""train"
" + ], + "text/plain": [ + "shape: (5, 4)\n", + "┌───────────┬──────────────────┬─────────────────────────┬───────┐\n", + "│ entity_id ┆ labour__birthday ┆ labour__native_language ┆ split │\n", + "│ --- ┆ --- ┆ --- ┆ --- │\n", + "│ str ┆ date ┆ str ┆ str │\n", + "╞═══════════╪══════════════════╪═════════════════════════╪═══════╡\n", + "│ E00001 ┆ 1954-11-11 ┆ spanish ┆ val │\n", + "│ E00002 ┆ 2004-04-16 ┆ japanese ┆ val │\n", + "│ E00003 ┆ 1950-08-17 ┆ english ┆ train │\n", + "│ E00004 ┆ 1992-06-23 ┆ russian ┆ val │\n", + "│ E00005 ┆ 1964-04-15 ┆ japanese ┆ train │\n", + "└───────────┴──────────────────┴─────────────────────────┴───────┘" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "from tab2seq.cohort import Cohort, CohortConfig, EntityInclusionCriteria\n", "\n", @@ -401,10 +648,66 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 12, "id": "07e6de56", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Vocabulary size: 114 tokens\n", + "Token columns: ['token_id', 'token', 'pretty_token', 'category', 'source_name', 'column_name', 'prefix', 'transform', 'count']\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "shape: (10, 9)
token_idtokenpretty_tokencategorysource_namecolumn_nameprefixtransformcount
i64strstrstrstrstrstrstri64
0"[PAD]""[PAD]""special""__special__""__special__""__special__""identity"-1
1"[UNK]""[UNK]""special""__special__""__special__""__special__""identity"-1
2"[CLS]""[CLS]""special""__special__""__special__""__special__""identity"-1
3"[SEP]""[SEP]""special""__special__""__special__""__special__""identity"-1
4"[MASK]""[MASK]""special""__special__""__special__""__special__""identity"-1
5"[DEATH]""[DEATH]""special""__special__""__special__""__special__""identity"-1
6"[RETIRED]""[RETIRED]""special""__special__""__special__""__special__""identity"-1
7"health__COST__BIN_0""COST__BIN_0""continuous_bin""health""cost""COST""continuous_bin"1408
8"health__COST__BIN_1""COST__BIN_1""continuous_bin""health""cost""COST""continuous_bin"1408
9"health__COST__BIN_10""COST__BIN_10""continuous_bin""health""cost""COST""continuous_bin"1409
" + ], + "text/plain": [ + "shape: (10, 9)\n", + "┌──────────┬────────────┬────────────┬────────────┬───┬────────────┬───────────┬───────────┬───────┐\n", + "│ token_id ┆ token ┆ pretty_tok ┆ category ┆ … ┆ column_nam ┆ prefix ┆ transform ┆ count │\n", + "│ --- ┆ --- ┆ en ┆ --- ┆ ┆ e ┆ --- ┆ --- ┆ --- │\n", + "│ i64 ┆ str ┆ --- ┆ str ┆ ┆ --- ┆ str ┆ str ┆ i64 │\n", + "│ ┆ ┆ str ┆ ┆ ┆ str ┆ ┆ ┆ │\n", + "╞══════════╪════════════╪════════════╪════════════╪═══╪════════════╪═══════════╪═══════════╪═══════╡\n", + "│ 0 ┆ [PAD] ┆ [PAD] ┆ special ┆ … ┆ __special_ ┆ __special ┆ identity ┆ -1 │\n", + "│ ┆ ┆ ┆ ┆ ┆ _ ┆ __ ┆ ┆ │\n", + "│ 1 ┆ [UNK] ┆ [UNK] ┆ special ┆ … ┆ __special_ ┆ __special ┆ identity ┆ -1 │\n", + "│ ┆ ┆ ┆ ┆ ┆ _ ┆ __ ┆ ┆ │\n", + "│ 2 ┆ [CLS] ┆ [CLS] ┆ special ┆ … ┆ __special_ ┆ __special ┆ identity ┆ -1 │\n", + "│ ┆ ┆ ┆ ┆ ┆ _ ┆ __ ┆ ┆ │\n", + "│ 3 ┆ [SEP] ┆ [SEP] ┆ special ┆ … ┆ __special_ ┆ __special ┆ identity ┆ -1 │\n", + "│ ┆ ┆ ┆ ┆ ┆ _ ┆ __ ┆ ┆ │\n", + "│ 4 ┆ [MASK] ┆ [MASK] ┆ special ┆ … ┆ __special_ ┆ __special ┆ identity ┆ -1 │\n", + "│ ┆ ┆ ┆ ┆ ┆ _ ┆ __ ┆ ┆ │\n", + "│ 5 ┆ [DEATH] ┆ [DEATH] ┆ special ┆ … ┆ __special_ ┆ __special ┆ identity ┆ -1 │\n", + "│ ┆ ┆ ┆ ┆ ┆ _ ┆ __ ┆ ┆ │\n", + "│ 6 ┆ [RETIRED] ┆ [RETIRED] ┆ special ┆ … ┆ __special_ ┆ __special ┆ identity ┆ -1 │\n", + "│ ┆ ┆ ┆ ┆ ┆ _ ┆ __ ┆ ┆ │\n", + "│ 7 ┆ health__CO ┆ COST__BIN_ ┆ continuous ┆ … ┆ cost ┆ COST ┆ continuou ┆ 1408 │\n", + "│ ┆ ST__BIN_0 ┆ 0 ┆ _bin ┆ ┆ ┆ ┆ s_bin ┆ │\n", + "│ 8 ┆ health__CO ┆ COST__BIN_ ┆ continuous ┆ … ┆ cost ┆ COST ┆ continuou ┆ 1408 │\n", + "│ ┆ ST__BIN_1 ┆ 1 ┆ _bin ┆ ┆ ┆ ┆ s_bin ┆ │\n", + "│ 9 ┆ health__CO ┆ COST__BIN_ ┆ continuous ┆ … ┆ cost ┆ COST ┆ continuou ┆ 1409 │\n", + "│ ┆ ST__BIN_10 ┆ 10 ┆ _bin ┆ ┆ ┆ ┆ s_bin ┆ │\n", + "└──────────┴────────────┴────────────┴────────────┴───┴────────────┴───────────┴───────────┴───────┘" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "from tab2seq.tokenization import Vocabulary, VocabularyConfig\n", "\n", @@ -437,10 +740,40 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 13, "id": "c428ab44", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "=== Token counts by source/category ===\n", + "shape: (4, 3)\n", + "┌─────────────┬────────────────┬─────┐\n", + "│ source_name ┆ category ┆ len │\n", + "│ --- ┆ --- ┆ --- │\n", + "│ str ┆ str ┆ u32 │\n", + "╞═════════════╪════════════════╪═════╡\n", + "│ health ┆ categorical ┆ 38 │\n", + "│ health ┆ continuous_bin ┆ 26 │\n", + "│ labour ┆ categorical ┆ 38 │\n", + "│ labour ┆ continuous_bin ┆ 5 │\n", + "└─────────────┴────────────────┴─────┘\n", + "\n", + "=== Column prefixes ===\n", + " health: {'cost': 'COST', 'procedure': 'PROC', 'length_of_stay': 'LOS', 'diagnosis': 'DIAG', 'department': 'DEPT'}\n", + " labour: {'native_language': 'LANG', 'weekly_hours': 'WEEKLY_HOURS', 'residence_region': 'REGION', 'occupation': 'OCC', 'status': 'STATUS'}\n", + "\n", + "=== COST bin edges (first 5) ===\n", + "[3.3400000e+00 9.3730000e+01 1.5992000e+02 2.3204000e+02 3.1103000e+02\n", + " 4.0185000e+02 5.0321000e+02 6.2021000e+02 7.5431000e+02 9.1167000e+02\n", + " 1.0985400e+03 1.3241600e+03 1.5917700e+03 1.9394900e+03 2.3920400e+03\n", + " 3.0152400e+03 3.8755600e+03 5.2860500e+03 7.6239200e+03 1.3011010e+04\n", + " 5.5285291e+05]\n" + ] + } + ], "source": [ "from tab2seq.tokenization import Tokenizer\n", "\n", @@ -466,14 +799,6 @@ "print(vocab.bin_edges_for(\"health\", \"cost\"))" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "0681faa6", - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "markdown", "id": "b83cf396", @@ -496,7 +821,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 14, "id": "bac090c0", "metadata": {}, "outputs": [], @@ -506,7 +831,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 15, "id": "bbef3537", "metadata": {}, "outputs": [], @@ -534,10 +859,19 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 16, "id": "cd361282", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Saved dataset as name: synthetic_demo_v1\n", + "Dataset dir: data/cohorts/synthetic_demo/datasets/4d582d55a824d090\n" + ] + } + ], "source": [ "dataset_name = \"synthetic_demo_v1\"\n", "artifacts = ds.write_parquet(dataset_name=dataset_name, force_write=True)\n", @@ -557,10 +891,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 17, "id": "62c8e718", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Loaded-by-name sample entity: E00882\n" + ] + } + ], "source": [ "# The registry lives under the cohort datasets directory by default.\n", "dataset_loaded = EventDataset.from_name(\n", @@ -599,10 +941,25 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 18, "id": "aeba953f", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "entity_id : E06941\n", + "split : train\n", + "static : {'entity_id': 'E06941', 'split': 'train', 'token_ids': [1], 'token_str': 'labour__native_language__swahili', 'labour__birthday': datetime.date(1986, 1, 30), 'labour__native_language': 'swahili'}\n", + "\n", + "First 3 events:\n", + " 2015-01-01 [labour] token_ids=[105, 86, 98, 110, 3] age=28\n", + " 2015-02-01 [labour] token_ids=[105, 86, 98, 111, 3] age=29\n", + " 2015-04-01 [labour] token_ids=[105, 86, 98, 113, 3] age=29\n" + ] + } + ], "source": [ "# iter_entity_records yields one dict per entity.\n", "# Each dict has keys: entity_id, split, static (dict), events (list of dicts).\n", @@ -630,10 +987,39 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 19, "id": "ed67f952", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "entity_id : E09428\n", + "static token_ids : [1]\n", + "\n", + "events DataFrame (64 rows):\n", + "shape: (64, 4)\n", + "┌───────────────────┬─────────────┬────────────────┬───────────┐\n", + "│ primary_timestamp ┆ source_name ┆ token_ids ┆ age_years │\n", + "│ --- ┆ --- ┆ --- ┆ --- │\n", + "│ str ┆ str ┆ list[i64] ┆ i64 │\n", + "╞═══════════════════╪═════════════╪════════════════╪═══════════╡\n", + "│ 2015-01-01 ┆ labour ┆ [102, 90, … 3] ┆ 38 │\n", + "│ 2015-03-01 ┆ labour ┆ [102, 90, … 3] ┆ 39 │\n", + "│ 2015-04-01 ┆ labour ┆ [102, 90, … 3] ┆ 39 │\n", + "│ 2015-05-01 ┆ labour ┆ [102, 90, … 3] ┆ 39 │\n", + "│ 2015-07-01 ┆ labour ┆ [102, 90, … 3] ┆ 39 │\n", + "│ … ┆ … ┆ … ┆ … │\n", + "│ 2024-04-01 ┆ labour ┆ [105, 83, … 3] ┆ 48 │\n", + "│ 2024-05-01 ┆ labour ┆ [105, 83, … 3] ┆ 48 │\n", + "│ 2024-08-01 ┆ labour ┆ [105, 83, … 3] ┆ 48 │\n", + "│ 2024-09-01 ┆ labour ┆ [105, 83, … 3] ┆ 48 │\n", + "│ 2024-11-01 ┆ labour ┆ [105, 83, … 3] ┆ 48 │\n", + "└───────────────────┴─────────────┴────────────────┴───────────┘\n" + ] + } + ], "source": [ "record_frame = dataset_loaded.sample_entity_record(\"train\", seed=7, format=\"frame\")\n", "\n", @@ -659,10 +1045,34 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 20, "id": "7cca8c49", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "entity_id : E09428\n", + "token_ids shape : (322,) (all events concatenated)\n", + "event_lengths : [5 5 5 5 5 5 5 5 5 5 5 5 6 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5\n", + " 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5] (tokens per event)\n", + "time : [16436 16495 16526 16556 16617 16648 16740 16832 16892 16953 16983 17075\n", + " 17152 17167 17226 17318 17348 17379 17410 17501 17532 17563 17652 17683\n", + " 17775 17805 17866 17897 17928 18017 18109 18170 18231 18322 18353 18444\n", + " 18506 18597 18628 18718 18748 18809 18901 18932 18962 19052 19083 19174\n", + " 19205 19297 19327 19358 19389 19448 19509 19570 19662 19692 19723 19814\n", + " 19844 19936 19967 20028] (days since reference per event)\n", + "temporal shape : (64, 2) [num_events, time + rel_date_features]\n", + "static_token_ids: [1]\n", + "\n", + "Per-event token lists (first 3 events):\n", + " event 0: [2, 102, 90, 99, 112]\n", + " event 1: [3, 102, 90, 99, 112]\n", + " event 2: [3, 102, 90, 99, 113]\n" + ] + } + ], "source": [ "import numpy as np\n", "\n", @@ -699,10 +1109,33 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 22, "id": "700def63", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "entity_id : E09428\n", + "token_ids shape : (65, 6) [num_events, max_event_len]\n", + "attention_mask shape: (65, 6)\n", + "time shape : (65,) [num_events]\n", + "\n", + "token_ids (first 4 events):\n", + "[[ 2 0 0 0 0 0]\n", + " [102 90 99 112 3 0]\n", + " [102 90 99 112 3 0]\n", + " [102 90 99 113 3 0]]\n", + "\n", + "attention_mask (first 4 events):\n", + "[[ True False False False False False]\n", + " [ True True True True True False]\n", + " [ True True True True True False]\n", + " [ True True True True True False]]\n" + ] + } + ], "source": [ "record_padded = dataset_loaded.sample_entity_record(\n", " \"train\", seed=7, format=\"padded_tensor\", pad_id=0\n", @@ -728,10 +1161,22 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 23, "id": "c36c712f", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "entity_id : E00003\n", + "n_events : 65\n", + "static : {'entity_id': 'E00003', 'split': 'train', 'labour__birthday': datetime.date(1950, 8, 17), 'labour__native_language': 'english'}\n", + "\n", + "missing entity → None\n" + ] + } + ], "source": [ "# Pick any entity ID that exists in the train split\n", "example_id = dataset_loaded.iter_entity_records(\"train\").__next__()[\"entity_id\"]\n", @@ -758,10 +1203,22 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 24, "id": "c8b49337", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Pulled 5 records from val split.\n", + "Entity IDs: ['E05348', 'E06753', 'E06327', 'E05762', 'E01476']\n", + "Event counts: [65, 70, 67, 62, 65]\n", + "\n", + "After reset, first entity: E06753 (same as E05348? False)\n" + ] + } + ], "source": [ "# Simulate a mini-batch loop: pull 5 entities, then reset and pull 3 more\n", "batch = []\n", @@ -780,19 +1237,11 @@ "r_reset = dataset_loaded.next_entity_record(\"val\", shuffle=True, seed=0)\n", "print(f\"\\nAfter reset, first entity: {r_reset['entity_id']} (same as {batch[0]['entity_id']}? {r_reset['entity_id'] == batch[0]['entity_id']})\")" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8c7f6963", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { "kernelspec": { - "display_name": "tab2seq (3.12.6)", + "display_name": "tab2seq (3.12.6.final.0)", "language": "python", "name": "python3" }, diff --git a/pyproject.toml b/pyproject.toml index 06c7dbc..118b1cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "tab2seq" -version = "0.1.9" +version = "0.1.10" description = "Transform tabular event data into sequences ready for Transformer and Sequential models: Life2Vec, BEHRT and more." readme = "README.md" requires-python = ">=3.11" diff --git a/src/tab2seq/source/config.py b/src/tab2seq/source/config.py index 972ae86..e8999b6 100644 --- a/src/tab2seq/source/config.py +++ b/src/tab2seq/source/config.py @@ -167,6 +167,21 @@ def _sync_output_folder_to_cache_dir(self) -> SourceConfig: self.cache_dir = Path(self.output_folder) return self + @model_validator(mode="after") + def _check_static_source_has_static_cols(self) -> SourceConfig: + if self.temporal_cols: + return self + all_feature_cols = (self.categorical_cols or []) + (self.continuous_cols or []) + non_static = [col.col_name for col in all_feature_cols if not col.static] + if non_static: + raise ValueError( + f"Source '{self.name}' has no temporal columns, so it is treated as a static " + f"source, but the following feature columns are not marked 'static=True': " + f"{non_static}. Either add temporal_cols to make it an event source, or set " + f"'static=True' on all feature columns." + ) + return self + @model_validator(mode="after") def _check_at_most_one_primary_timestamp(self) -> SourceConfig: if not self.temporal_cols: @@ -233,5 +248,4 @@ def primary_temporal(self) -> TemporalColConfig | None: """The primary temporal column config, if any.""" if not self.temporal_cols: return None - return next((col for col in self.temporal_cols if col.is_primary), None) - + return next((col for col in self.temporal_cols if col.is_primary), None) \ No newline at end of file