diff --git a/scripts/plots.py b/scripts/plots.py index 398cae0..bf810d8 100644 --- a/scripts/plots.py +++ b/scripts/plots.py @@ -16,7 +16,7 @@ from impact.cli import inject_impact_config from impact.config import ImpactConfig from impact.logger import configure_logging -from impact.pipelines import DataPipeline, IMPACTEvaluationPipeline +from impact.pipelines import DataPipeline, ImpactEvaluationPipeline from impact.utils import seed_everything logger = structlog.stdlib.get_logger(__name__) @@ -154,20 +154,10 @@ def cli() -> None: @click.option( "--lang", type=click.Choice(["pl", "en"], case_sensitive=False), - default="pl", + default="en", show_default=True, help="Language for axis labels and annotations.", ) -@click.option( - "--n-partitions", - type=int, - default=1, - show_default=True, - help=""" - Number of partitions for sequential processing. Use >1 for large datasets that - exceed available RAM (data is partitioned by user, i.e. uid column). - """, -) @click.option( "--raw-data-dir", type=click.Path(file_okay=False, path_type=Path), @@ -189,7 +179,7 @@ def cli() -> None: @click.option( "--processed-data-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Base directory for processed data (the final, canonical datasets for modeling). Defaults to /data/processed. @@ -269,9 +259,8 @@ def cli() -> None: def plot_alpha_curve( impact_config: ImpactConfig, sasrec_checkpoint_path: Path, - log_level: str, - lang: str, - n_partitions: int, + log_level: str = "WARNING", + lang: str = "en", raw_data_dir: Path | None = None, interim_data_dir: Path | None = None, processed_data_dir: Path | None = None, @@ -320,11 +309,11 @@ def plot_alpha_curve( config=impact_config, ) - logger.info("Fetching raw events data...") - events_lf = data_pipeline.get_raw_events() + logger.info("Fetching listening events data...") + listens_lf = data_pipeline.get_listen_events() logger.info("Fetching train/val/test splits...") - splits, item_index = data_pipeline.get_splits(n_partitions=n_partitions) + splits, item_index = data_pipeline.get_splits() logger.info("Tensorizing splits...") tensorized_splits = data_pipeline.tensorize(splits) @@ -346,7 +335,7 @@ def plot_alpha_curve( ) logger.info(f"Running IMPACT evaluation for {alpha=}...") - eval_pipeline = IMPACTEvaluationPipeline(eval_config, item_index, events_lf) + eval_pipeline = ImpactEvaluationPipeline(eval_config, item_index, listens_lf) metrics: dict[str, float] = eval_pipeline.evaluate( split_partition=tensorized_splits.val, sasrec_state_dict_path=sasrec_checkpoint_path, @@ -521,7 +510,7 @@ def _create_subplot( @click.option( "--lang", type=click.Choice(["pl", "en"], case_sensitive=False), - default="pl", + default="en", show_default=True, help="Language for axis labels and annotations.", ) @@ -539,9 +528,9 @@ def _create_subplot( "--format", "-f", "formats", - type=click.Choice(["svg", "pdf", "png"], case_sensitive=False), + type=click.Choice(("svg", "pdf", "png"), case_sensitive=False), multiple=True, - default=["svg", "pdf"], + default=("svg", "pdf"), show_default=True, help="Output format(s). Can be specified multiple times: -f svg -f pdf.", ) @@ -554,11 +543,11 @@ def _create_subplot( ) def plot_sasrec_training( log_file_path: Path, - log_level: str, - lang: str, - out_dir: Path | None, - formats: list[Literal["svg", "pdf", "png"]], - png_scale: float, + log_level: str = "INFO", + lang: str = "en", + out_dir: Path | None = None, + formats: tuple[Literal["svg", "pdf", "png"], ...] = ("svg", "pdf"), + png_scale: float = 3.0, ) -> None: """Plot SASRec training curves from a structured NDJSON log file. diff --git a/scripts/tables.py b/scripts/tables.py index 3ca535c..4bd2a64 100644 --- a/scripts/tables.py +++ b/scripts/tables.py @@ -37,14 +37,14 @@ def cli() -> None: @click.option( "--lang", type=click.Choice(["pl", "en"], case_sensitive=False), - default="pl", + default="en", help="Language for the table headers and textual content (pl or en).", ) @click.option( "--out-dir", "-o", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Output directory for the LaTeX code and corresponding PipelineConfig YAML dumps. Defaults to project root / 'reports' / 'tables' / 'data-funnel-stats'. @@ -53,7 +53,7 @@ def cli() -> None: @click.option( "--raw-data-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Root directory for raw data (the original, immutable data dump). Defaults to project root / 'data' / 'raw'. @@ -62,7 +62,7 @@ def cli() -> None: @click.option( "--interim-data-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Root directory for interim data (intermediate data that has been transformed). Defaults to project root / 'data' / 'interim'. @@ -71,7 +71,7 @@ def cli() -> None: @click.option( "--processed-data-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Root directory for processed data (the final, canonical datasets for modeling). Defaults to project root / 'data' / 'processed'. @@ -80,8 +80,8 @@ def cli() -> None: @inject_impact_config def generate_data_funnel_stats_table( impact_config: ImpactConfig, - log_level: str, - lang: str = "pl", + log_level: str = "INFO", + lang: str = "en", out_dir: Path | None = None, raw_data_dir: Path | None = None, interim_data_dir: Path | None = None, @@ -134,20 +134,17 @@ def generate_data_funnel_stats_table( class _VariantConfig(NamedTuple): step: int - n_partitions: int dataset_name: str # Name of the dataset for the "reporting" (in the thesis) marker: str # Marker for the footnotes variants: dict[DatasetSize, _VariantConfig] = { DatasetSize.SIZE_50M: _VariantConfig( step=5, - n_partitions=1, dataset_name="Yambda-50M", marker="a", ), DatasetSize.SIZE_5B: _VariantConfig( step=20, - n_partitions=32, dataset_name="Yambda-5B", marker="b", ), @@ -166,12 +163,13 @@ class _VariantConfig(NamedTuple): pipeline.config = variant_config # Switch pipeline config events_lf = pipeline.get_raw_events() - sessions_lf = pipeline.get_sessions(n_partitions=variant.n_partitions) - sequences_lf = pipeline.get_sequences(n_partitions=variant.n_partitions) + sessions_lf = pipeline.get_sessions() + sequences_lf = pipeline.get_sequences() tnote = f"\\tnote{{\\hyperlink{{note:{variant.marker}}}{{{variant.marker}}}}}" rows.append({ "dataset": f"{variant.dataset_name}{tnote}", + # TODO: report also the number of valid listening events "raw": _format_num(events_lf.select(pl.len()).collect().item()), "sessionized": _format_num(sessions_lf.select(pl.len()).collect().item()), "sequenced": _format_num(sequences_lf.select(pl.len()).collect().item()), diff --git a/src/impact/cli.py b/src/impact/cli.py index 46fdf79..837f7a6 100644 --- a/src/impact/cli.py +++ b/src/impact/cli.py @@ -57,15 +57,6 @@ def inject_impact_config(f: Callable[..., Any]) -> Callable[..., Any]: variants contain only a specific subset. """, ), - click.option( - "--session-inactivity-gap-seconds", - type=int, - default=1200, - show_default=True, - help=""" - Number of seconds of inactivity before a session is considered finished. - """, - ), click.option( "--session-tracks-per-session", type=int, @@ -77,17 +68,6 @@ def inject_impact_config(f: Callable[..., Any]) -> Callable[..., Any]: K tracks are truncated. """, ), - click.option( - "--session-min-listen-seconds", - type=int, - default=30, - show_default=True, - help=""" - Minimum playback duration (in seconds) to consider an event as a valid - listen. Defaults to 30 (industry standard used by authors of other - publications). - """, - ), click.option( "--sequence-context-sessions", type=int, diff --git a/src/impact/config.py b/src/impact/config.py index bc75991..c3e9b2f 100644 --- a/src/impact/config.py +++ b/src/impact/config.py @@ -39,9 +39,6 @@ def __post_init__(self) -> None: class SessionConfig: """Configuration for the sessionization stage (Algorithm 1 in the thesis).""" - inactivity_gap_seconds: int = 1200 - """Number of seconds of inactivity before a session is considered finished.""" - tracks_per_session: int = 10 """Number of tracks that constitute a session. Denoted as K in the thesis. @@ -49,12 +46,6 @@ class SessionConfig: are truncated. """ - min_listen_seconds: int = 30 - """Minimum playback duration (in seconds) to consider an event as a valid listen. - - Defaults to 30 (industry standard used by authors of other publications). - """ - @dataclass(frozen=True) class SequenceConfig: @@ -241,9 +232,7 @@ def sessionization_id(self) -> str: """Cache key for the sessionization stage (Algorithm 1 in the thesis).""" parts = [ "sessions", - f"gap{self.session.inactivity_gap_seconds}", f"k{self.session.tracks_per_session}", - f"min{self.session.min_listen_seconds}", ] return "_".join(parts) diff --git a/src/impact/features/__init__.py b/src/impact/features/__init__.py index 5abacf0..73e6f9c 100644 --- a/src/impact/features/__init__.py +++ b/src/impact/features/__init__.py @@ -1,3 +1,4 @@ +from .preprocessing import filter_listen_events from .sequence_extraction import extract_sequences from .sessionization import sessionize from .temporal_split import split @@ -5,6 +6,7 @@ __all__ = [ "extract_sequences", + "filter_listen_events", "sessionize", "split", "tensorize", diff --git a/src/impact/features/preprocessing.py b/src/impact/features/preprocessing.py new file mode 100644 index 0000000..dd24f6c --- /dev/null +++ b/src/impact/features/preprocessing.py @@ -0,0 +1,54 @@ +"""Raw event preprocessing applied before any feature-engineering step.""" + +import polars as pl +import structlog + +logger = structlog.stdlib.get_logger(__name__) + +MIN_LISTEN_SECONDS: int = 30 +"""Minimum playback duration (in seconds) to consider an event as a valid listen. +Below this threshold events are treated as noise (accidental plays, skips). +Follows the industry standard established in the literature. +""" + + +def filter_listen_events( + events_lf: pl.LazyFrame, + min_listen_seconds: int = MIN_LISTEN_SECONDS, +) -> pl.LazyFrame: + """Filter raw events of all types down to valid listening events only. + + This is the only place that decides what constitutes a valid listen, so the + definition stays consistent throughout the entire codebase. + + Args: + events_lf: Raw LazyFrame containing all event types. + min_listen_seconds: Minimum playback duration (in seconds) to consider an event + as a valid listen. Below this threshold events are treated as noise + (accidental plays, skips). + + Returns: + A LazyFrame containing only `event_type == "listen"` rows where the user + actually played enough of the track, with an `actual_played_seconds` column + added. + """ + logger.debug( + "Filtering raw events to valid listen events", + min_listen_seconds=min_listen_seconds, + ) + return ( + events_lf + + # Filter out non-listen events + .filter(pl.col("event_type") == "listen") + + # Calculate actual played time in seconds + .with_columns( + (pl.col("track_length_seconds") * pl.col("played_ratio_pct") / 100).alias( + "actual_played_seconds" + ) + ) + + # Filter valid listening events + .filter(pl.col("actual_played_seconds") >= min_listen_seconds) + ) # fmt: skip diff --git a/src/impact/features/sessionization.py b/src/impact/features/sessionization.py index 2309f74..3589b95 100644 --- a/src/impact/features/sessionization.py +++ b/src/impact/features/sessionization.py @@ -1,54 +1,42 @@ -"""Module for processing raw events from the dataset and extracting sessions.""" - import polars as pl import structlog logger = structlog.stdlib.get_logger(__name__) +INACTIVITY_GAP_SECONDS: int = 1200 +"""Number of seconds of inactivity before a session is considered finished. +Follows the industry standard used across the music streaming literature. +""" + def sessionize( - events_lf: pl.LazyFrame, - inactivity_gap_seconds: int = 1200, + listens_lf: pl.LazyFrame, tracks_per_session: int = 10, # K - min_listen_seconds: int = 30, # Threshold defining a valid listening event + inactivity_gap_seconds: int = INACTIVITY_GAP_SECONDS, ) -> pl.LazyFrame: - """Process raw data containing all event types and extract user-scoped sessions. + """Extract user-scoped sessions from pre-filtered listen events. Args: - events_lf: A Polars LazyFrame containing all event types. - inactivity_gap_seconds: Number of seconds of inactivity before a session is - considered finished. Defaults to 1200 (20 minutes). - tracks_per_session: Number of tracks that consitute a session. Denoted as K in + listens_lf: LazyFrame of valid listen events (event_type == "listen", + actual_played_seconds already computed and filtered). + tracks_per_session: Number of tracks that constitute a session. Denoted as K in the thesis. Sessions with less than K tracks are discarded. Sessions with more than K tracks are truncated. Defaults to 10. - min_listen_seconds: Minimum playback duration (in seconds) to consider an event - as a valid listen. Defaults to 30 (industry standard used by authors of - other publications). + inactivity_gap_seconds: Number of seconds of inactivity before a session is + considered finished. Returns: - A Polars LazyFrame extending the events_lf by the session_id column containing + A Polars LazyFrame extending the listens_lf by the session_id column containing assigned session IDs, unique per user (not globally unique). """ logger.info( - "Preprocessing raw events into sessions", + "Extracting user-scoped sessions from listen events", inactivity_gap_seconds=inactivity_gap_seconds, - k_tracks=tracks_per_session, + tracks_per_session=tracks_per_session, ) sessions_lf = ( - events_lf - # Filter out non-listen events - .filter(pl.col("event_type") == "listen") - - # Calculate actual played time in seconds - .with_columns( - ( - pl.col("track_length_seconds") * pl.col("played_ratio_pct") / 100 - ).alias("actual_played_seconds") - ) - - # Filter valid listening events - .filter(pl.col("actual_played_seconds") >= min_listen_seconds) + listens_lf # Sort temporally .sort(["uid", "timestamp"]) diff --git a/src/impact/main.py b/src/impact/main.py index d93da69..b0848a9 100644 --- a/src/impact/main.py +++ b/src/impact/main.py @@ -6,7 +6,7 @@ import json import logging from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import click import dotenv @@ -17,9 +17,10 @@ from impact.logger import configure_logging from impact.pipelines import ( DataPipeline, - IMPACTEvaluationPipeline, + ImpactEvaluationPipeline, SASRecTrainingPipeline, ) +from impact.schemas import TensorizedDataset from impact.utils import seed_everything if TYPE_CHECKING: @@ -46,16 +47,6 @@ def cli() -> None: show_default=True, help="Logging verbosity level.", ) -@click.option( - "--n-partitions", - type=int, - default=1, - show_default=True, - help=""" - Number of partitions for sequential processing. Use >1 for large datasets that - exceed available RAM (data is partitioned by user, i.e. uid column). - """, -) @click.option( "--raw-data-dir", type=click.Path(file_okay=False, path_type=Path), @@ -77,7 +68,7 @@ def cli() -> None: @click.option( "--processed-data-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Base directory for processed data (the final, canonical datasets for modeling). Defaults to /data/processed. @@ -86,13 +77,13 @@ def cli() -> None: @click.option( "--base-logs-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help="Base directory for log files. Defaults to /logs.", ) @click.option( "--models-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Base directory where a serialized checkpoint of the trained SASRecNBR model will be saved. Defaults to /models. @@ -101,8 +92,7 @@ def cli() -> None: @inject_impact_config def train_sasrec( impact_config: ImpactConfig, - log_level: str, - n_partitions: int, + log_level: str = "INFO", raw_data_dir: Path | None = None, interim_data_dir: Path | None = None, processed_data_dir: Path | None = None, @@ -155,7 +145,7 @@ def train_sasrec( ) logger.info("Fetching train/val/test splits...") - splits, item_index = data_pipeline.get_splits(n_partitions=n_partitions) + splits, item_index = data_pipeline.get_splits() logger.info("Tensorizing splits...") tensorized_splits = data_pipeline.tensorize(splits) @@ -184,10 +174,20 @@ def train_sasrec( @cli.command(name="evaluate-impact") @click.option( "--sasrec-checkpoint-path", - type=click.Path(dir_okay=False, path_type=Path), + type=click.Path(exists=True, dir_okay=False, path_type=Path), required=True, help="Path to a trained SASRecNBR checkpoint that will be evaluated.", ) +@click.option( + "--partition", + type=click.Choice(["val", "test"], case_sensitive=False), + default="test", + show_default=True, + help=""" + Dataset partition to evaluate on. Use 'val' during hyperparameter search, + 'test' for final evaluation on the held-out data. + """, +) @click.option( "--log-level", type=click.Choice(list(logging.getLevelNamesMapping()), case_sensitive=False), @@ -195,16 +195,6 @@ def train_sasrec( show_default=True, help="Logging verbosity level.", ) -@click.option( - "--n-partitions", - type=int, - default=1, - show_default=True, - help=""" - Number of partitions for sequential processing. Use >1 for large datasets that - exceed available RAM (data is partitioned by user, i.e. uid column). - """, -) @click.option( "--raw-data-dir", type=click.Path(file_okay=False, path_type=Path), @@ -226,7 +216,7 @@ def train_sasrec( @click.option( "--processed-data-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Base directory for processed data (the final, canonical datasets for modeling). Defaults to /data/processed. @@ -235,33 +225,28 @@ def train_sasrec( @click.option( "--base-logs-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help="Base directory for log files. Defaults to /logs.", ) @inject_impact_config def evaluate_impact( impact_config: ImpactConfig, sasrec_checkpoint_path: Path, - log_level: str, - n_partitions: int, + partition: str = "test", + log_level: str = "INFO", raw_data_dir: Path | None = None, interim_data_dir: Path | None = None, processed_data_dir: Path | None = None, base_logs_dir: Path | None = None, exec_time: datetime.datetime | None = None, ) -> None: - """Run IMPACT late-fusion evaluation on the held-out test split. + """Run IMPACT late-fusion evaluation on the specified split. Loads a trained SASRecNBR checkpoint, builds the ACT-R history index for test users, runs batched inference, and computes sample-weighted metrics at all configured cutoff levels: global, repetition-focused, and exploration-focused NDCG@K and Recall@K. """ - if not sasrec_checkpoint_path.exists(): - raise FileNotFoundError( - f"SASRecNBR checkpoint not found: {sasrec_checkpoint_path.as_posix()}" - ) - timestamp = exec_time or datetime.datetime.now(datetime.UTC) raw_data_dir = raw_data_dir or PROJECT_ROOT / "data" / "raw" @@ -299,19 +284,24 @@ def evaluate_impact( config=impact_config, ) - logger.info("Fetching raw events...") - events_lf = data_pipeline.get_raw_events() + logger.info("Fetching listen events...") + listens_lf = data_pipeline.get_listen_events() logger.info("Fetching train/val/test splits...") - splits, item_index = data_pipeline.get_splits(n_partitions=n_partitions) + splits, item_index = data_pipeline.get_splits() logger.info("Tensorizing splits...") tensorized_splits = data_pipeline.tensorize(splits) - evaluation_pipeline = IMPACTEvaluationPipeline(impact_config, item_index, events_lf) + evaluation_pipeline = ImpactEvaluationPipeline( + config=impact_config, + item_index=item_index, + listens_lf=listens_lf, + ) logger.info("Running IMPACT evaluation...") - metrics = evaluation_pipeline.evaluate(tensorized_splits, sasrec_checkpoint_path) + split_partition = cast(TensorizedDataset, getattr(tensorized_splits, partition)) + metrics = evaluation_pipeline.evaluate(split_partition, sasrec_checkpoint_path) logger.info("IMPACT evaluation complete", **metrics) @@ -323,16 +313,6 @@ def evaluate_impact( show_default=True, help="Logging verbosity level.", ) -@click.option( - "--n-partitions", - type=int, - default=1, - show_default=True, - help=""" - Number of partitions for sequential processing. Use >1 for large datasets that - exceed available RAM (data is partitioned by user, i.e. uid column). - """, -) @click.option( "--raw-data-dir", type=click.Path(file_okay=False, path_type=Path), @@ -354,7 +334,7 @@ def evaluate_impact( @click.option( "--processed-data-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Base directory for processed data (the final, canonical datasets for modeling). Defaults to /data/processed. @@ -363,13 +343,13 @@ def evaluate_impact( @click.option( "--base-logs-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help="Base directory for log files. Defaults to /logs.", ) @click.option( "--models-dir", type=click.Path(file_okay=False, path_type=Path), - required=False, + default=None, help=""" Base directory where the grid search artifacts will be saved: the best-performing SASRecNBR model checkpoint, optimal ImpactConfig snapshot, @@ -379,8 +359,7 @@ def evaluate_impact( @inject_impact_config def grid_search( impact_config: ImpactConfig, - log_level: str, - n_partitions: int, + log_level: str = "INFO", raw_data_dir: Path | None = None, interim_data_dir: Path | None = None, processed_data_dir: Path | None = None, @@ -434,7 +413,7 @@ def grid_search( ) logger.info("Fetching train/val/test splits...") - splits, item_index = data_pipeline.get_splits(n_partitions=n_partitions) + splits, item_index = data_pipeline.get_splits() logger.info("Tensorizing splits...") tensorized_splits = data_pipeline.tensorize(splits) diff --git a/src/impact/models/actr/history_index.py b/src/impact/models/actr/history_index.py index 37d9cce..0e1ba57 100644 --- a/src/impact/models/actr/history_index.py +++ b/src/impact/models/actr/history_index.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, final import numpy as np import polars as pl @@ -14,6 +14,7 @@ logger = structlog.stdlib.get_logger(__name__) +@final class ACTRHistoryIndex: """In-memory index of user event histories for ACT-R scoring. @@ -24,7 +25,7 @@ class ACTRHistoryIndex: def __init__( self, - events_lf: pl.LazyFrame, + listens_lf: pl.LazyFrame, eval_uids: list[int] | pl.Series, item_index: ItemIndex, ) -> None: @@ -35,7 +36,9 @@ def __init__( timestamp for efficient cutoff slicing. Args: - events_lf: LazyFrame of all interaction events. + listens_lf: LazyFrame of valid listen events (event_type == "listen", + actual_played_seconds already computed and filtered). Contains the + listening history for all users. eval_uids: User IDs in the evaluation split; only these are indexed. item_index: Item-to-index mapping for ID translation derived from the training split. @@ -49,11 +52,11 @@ def __init__( self._history: dict[int, dict[str, np.ndarray]] = {} for uid, df in tqdm( iterable=( - events_lf.filter( - pl.col("event_type") == "listen", - pl.col("uid").is_in(eval_uids), - ) + # Filter out non-eval users + listens_lf.filter(pl.col("uid").is_in(eval_uids)) + # Map item IDs to indices .join(item_index.table.lf, on="item_id", how="inner") + # The remaining operations are obvious .select(["uid", "idx", "timestamp", "is_organic"]) .collect() .partition_by("uid", as_dict=True) diff --git a/src/impact/pipelines/__init__.py b/src/impact/pipelines/__init__.py index 5f87dad..4039fb8 100644 --- a/src/impact/pipelines/__init__.py +++ b/src/impact/pipelines/__init__.py @@ -1,9 +1,9 @@ from .data_pipeline import DataPipeline -from .impact_evaluation_pipeline import IMPACTEvaluationPipeline +from .impact_evaluation_pipeline import ImpactEvaluationPipeline from .sasrec_training_pipeline import SASRecTrainingPipeline __all__ = [ "DataPipeline", - "IMPACTEvaluationPipeline", + "ImpactEvaluationPipeline", "SASRecTrainingPipeline", ] diff --git a/src/impact/pipelines/data_pipeline.py b/src/impact/pipelines/data_pipeline.py index fdd583c..530d1a9 100644 --- a/src/impact/pipelines/data_pipeline.py +++ b/src/impact/pipelines/data_pipeline.py @@ -3,8 +3,17 @@ import polars as pl import structlog +from tqdm import tqdm from impact.config import ImpactConfig +from impact.dataset import YambdaFetcher +from impact.features import ( + extract_sequences, + filter_listen_events, + sessionize, + split, + tensorize, +) from impact.schemas import ( ItemIndex, ParquetTableSet, @@ -12,6 +21,7 @@ TensorizedDataset, TensorizedSplit, ) +from impact.utils import estimate_processing_partitions logger = structlog.stdlib.get_logger(__name__) @@ -19,10 +29,18 @@ class DataPipeline: """Orchestrates IMPACT data processing with a cache-aware API. - This pipeline manages the dataset acquisition and transformation of raw events into - ML-ready temporal splits. It ensures reproducibility and avoids redundant - computations by caching intermediate results in nested directories based on the - pipeline configuration. + Manages dataset acquisition and transformation of raw events into ML-ready tensors. + Intermediate results are cached in nested directories keyed by IMPACT Framework + configuration, ensuring reproducibility and avoiding redundant computation across + runs. + + Processing stages in order: + 1. `get_raw_events()` - fetch and cache raw Yambda events (multi-event type). + 2. `get_listen_events()` - filter to valid listen events only. + 3. `get_sessions()` - extract user-scoped listening sessions. + 4. `get_sequences()` - build context/target sequence pairs. + 5. `get_splits()` - apply temporal train/val/test split. + 6. `tensorize()` - convert to memory-mapped NumPy arrays. """ def __init__( @@ -48,9 +66,15 @@ def __init__( self.config: ImpactConfig = config def get_raw_events(self) -> pl.LazyFrame: - """Fetch raw interactions.""" - from impact.dataset import YambdaFetcher + """Get raw interaction events, downloading and caching if not already cached. + + Caching is handled by the Hugging Face Hub client, which persists the dataset + under `` and skips downloading on subsequent calls. + Returns: + A LazyFrame of raw interaction events of all types, scanned lazily from the + cached arquet files. + """ fetcher = YambdaFetcher( cache_dir=self.raw_dir, dataset_type=self.config.dataset.format, @@ -58,155 +82,189 @@ def get_raw_events(self) -> pl.LazyFrame: ) return fetcher.get_interactions(event_type=self.config.dataset.interaction_type) - def get_sessions(self, n_partitions: int = 1) -> pl.LazyFrame: - """Get sessionized data, computing it if not cached. + def get_listen_events(self) -> pl.LazyFrame: + """Get valid listen events, computing and caching to disk if not already cached. - Args: - n_partitions: Number of partitions for sequential processing. Use >1 for - large datasets that exceed available RAM (data is partitioned by user, - i.e. uid column). + Filters raw events down to the listening events that meet the minimum playback + threshold defined in config. Results are persisted under + `//listen_events/` and reused on subsequent calls. + + Returns: + A LazyFrame of valid listen events with an `actual_played_seconds` column + added. Scanned lazily from the cached parquet file. """ - if n_partitions < 1: - raise ValueError("Number of partitions must be greater than 0.") + cache_dir = self.interim_dir / self.config.dataset_id / "listen_events" + success_marker = cache_dir / "_SUCCESS" + + # Atomicity check + if not success_marker.exists(): + # Clean up potentially corrupted partial runs + if cache_dir.exists(): + shutil.rmtree(cache_dir) + cache_dir.mkdir(parents=True) + + logger.info("No complete cached listen events found, computing...") + + with tqdm( + desc="Computing listen events", + unit="", + bar_format="{desc}: {elapsed} {postfix}", + ) as bar: + filter_listen_events(events_lf=self.get_raw_events()).sink_parquet( + path=cache_dir / "listen_events.parquet", + ) + bar.set_postfix_str("done") + + success_marker.touch() + logger.info("Listen events computed and cached.") - from impact.features import sessionize + return pl.scan_parquet(cache_dir / "listen_events.parquet") - sessions_dir = ( + def get_sessions(self) -> pl.LazyFrame: + """Get user-scoped sessions, computing and caching if not already cached. + + Partitions the listen events by user (`uid % n_partitions`) and processes each + partition sequentially to fit within available RAM. The number of partitions + is estimated automatically based on dataset size. + + Results are persisted under `///` + and reused on subsequent calls. + + Returns: + A LazyFrame of extracted user-scoped listening sessions, scanned lazily from + the cached parquet files. Each row is a listen event annotated with its + `session_id` (it's unique per user, not globally). + """ + cache_dir = ( self.interim_dir / self.config.dataset_id / self.config.sessionization_id ) - success_marker = sessions_dir / "_SUCCESS" + success_marker = cache_dir / "_SUCCESS" # Atomicity check if not success_marker.exists(): + # Clean up potentially corrupted partial runs + if cache_dir.exists(): + shutil.rmtree(cache_dir) + cache_dir.mkdir(parents=True) + + listens_lf = self.get_listen_events() + n_partitions = estimate_processing_partitions(data=listens_lf) logger.info( "No complete sessionized data found in cache, computing...", partitions=n_partitions, ) - # Clean up potentially corrupted partial runs - if sessions_dir.exists(): - shutil.rmtree(sessions_dir) - - events_lf = self.get_raw_events() - - sessions_dir.mkdir(parents=True) - - for i in range(n_partitions): + for i in tqdm(range(n_partitions), desc="Sessionizing", unit="partition"): logger.info("Computing partition", partition=i + 1, total=n_partitions) sessionize( - events_lf=events_lf.filter(pl.col("uid") % n_partitions == i), - inactivity_gap_seconds=self.config.session.inactivity_gap_seconds, + listens_lf=listens_lf.filter(pl.col("uid") % n_partitions == i), tracks_per_session=self.config.session.tracks_per_session, - min_listen_seconds=self.config.session.min_listen_seconds, - ).sink_parquet(sessions_dir / f"sessions_part_{i}.parquet") + ).sink_parquet(cache_dir / f"sessions_part_{i}.parquet") # Mark as completed success_marker.touch() - return pl.scan_parquet(sessions_dir / "sessions_part_*.parquet") + return pl.scan_parquet(cache_dir / "sessions_part_*.parquet") - def get_sequences(self, n_partitions: int = 1) -> pl.LazyFrame: - """Get extracted sequences of sessions, computing them if not cached. + def get_sequences(self) -> pl.LazyFrame: + """Get extracted sequences of sessions, computing if not already cached. - Args: - n_partitions: Number of partitions for sequential processing. Use >1 for - large datasets that exceed available RAM (data is partitioned by user, - i.e. uid column). - """ - if n_partitions < 1: - raise ValueError("Number of partitions must be greater than 0.") + Partitions the sessions by user (`uid % n_partitions`) and processes each + partition sequentially to fit within available RAM. The number of partitions + is estimated automatically based on dataset size. - from impact.features import extract_sequences + Results are persisted under + `////` + and reused on subsequent calls. - sequences_dir = ( + Returns: + A LazyFrame of sequences scanned lazily from the cached parquet files, + with a globally unique `sequence_idx` column assigned across all partitions + (1-based). + """ + cache_dir = ( self.interim_dir / self.config.dataset_id / self.config.sessionization_id / self.config.sequence_extraction_id ) - success_marker = sequences_dir / "_SUCCESS" + success_marker = cache_dir / "_SUCCESS" # Atomicity check if not success_marker.exists(): + # Clean up potentially corrupted partial runs + if cache_dir.exists(): + shutil.rmtree(cache_dir) + cache_dir.mkdir(parents=True) + + sessions_lf = self.get_sessions() + n_partitions = estimate_processing_partitions(data=sessions_lf) logger.info( "No complete sequenced data found in cache, computing...", partitions=n_partitions, ) - # Clean up potentially corrupted partial runs - if sequences_dir.exists(): - shutil.rmtree(sequences_dir) - - sessions_lf = self.get_sessions(n_partitions=n_partitions) - - sequences_dir.mkdir(parents=True) - - for i in range(n_partitions): + for i in tqdm(range(n_partitions), desc="Sequencing", unit="partition"): logger.info("Computing partition", partition=i + 1, total=n_partitions) - extract_sequences( sessions_lf=sessions_lf.filter(pl.col("uid") % n_partitions == i), context_sessions=self.config.sequence.context_sessions, step=self.config.sequence.step, assign_index=False, # Assign global index after merging - ).sink_parquet(sequences_dir / f"sequences_part_{i}.parquet") + ).sink_parquet(cache_dir / f"sequences_part_{i}.parquet") # Mark as completed success_marker.touch() # Assign global index - return pl.scan_parquet( - sequences_dir / "sequences_part_*.parquet" - ).with_row_index("sequence_idx", offset=1) + return pl.scan_parquet(cache_dir / "sequences_part_*.parquet").with_row_index( + "sequence_idx", offset=1 + ) - def get_splits( - self, - n_partitions: int = 1, - compute_ram_fraction: float = 0.15, - ) -> tuple[TemporalSplit, ItemIndex]: - """Get final train/val/test splits, computing them if not cached. + def get_splits(self) -> tuple[TemporalSplit, ItemIndex]: + """Get temporal train/val/test splits and item index, computing if not cached. - Args: - n_partitions: Number of partitions for sequential processing. Use >1 for - large datasets that exceed available RAM (data is partitioned by user, - i.e. uid column). - compute_ram_fraction: Fraction of available RAM allowed per processing - partition. It should be kept low (<=0.15) to account for memory spikes - during joins, explodes, and struct (non-streamable) Polars operations. - """ - if n_partitions < 1: - raise ValueError("Number of partitions must be greater than 0.") + Unlike other pipeline stages, partitioning and disk writes are handled + internally by `impact.features.split()`. This is intentional because the + item-to-index mapping is derived from the train set and must be known before + val/test sets can be written, making the entire operation atomic. Partition + count and RAM management are estimated automatically inside `split()`. - from impact.features import split + Results are persisted under `/// + //` and reused on subsequent calls. - splits_dir = ( + Returns: + A tuple of (TemporalSplit, ItemIndex) where splits are scanned lazily + from cached parquet files and the item index maps train-set item IDs to + contiguous 1-based indices (unseen items map to 0 which serves as padding). + """ + cache_dir = ( self.processed_dir / self.config.dataset_id / self.config.sessionization_id / self.config.sequence_extraction_id / self.config.temporal_split_id ) - item_to_idx_cache_dir = splits_dir / "item_to_idx" - train_cache_dir = splits_dir / "train" - val_cache_dir = splits_dir / "val" - test_cache_dir = splits_dir / "test" - success_marker = splits_dir / "_SUCCESS" + item_to_idx_cache_dir = cache_dir / "item_to_idx" + train_cache_dir = cache_dir / "train" + val_cache_dir = cache_dir / "val" + test_cache_dir = cache_dir / "test" + success_marker = cache_dir / "_SUCCESS" # Atomicity check if not success_marker.exists(): - logger.info("No complete temporal splits found in cache, computing...") - # Clean up potentially corrupted partial runs - if splits_dir.exists(): - shutil.rmtree(splits_dir) - - sequences_lf = self.get_sequences(n_partitions=n_partitions) - + if cache_dir.exists(): + shutil.rmtree(cache_dir) + cache_dir.mkdir(parents=True) item_to_idx_cache_dir.mkdir(parents=True) train_cache_dir.mkdir(parents=True) val_cache_dir.mkdir(parents=True) test_cache_dir.mkdir(parents=True) + sequences_lf = self.get_sequences() + logger.info("No complete temporal splits found in cache, computing...") + # Discard the result as it has to be recreated at the end either way # (because of caching) _ = split( @@ -218,7 +276,6 @@ def get_splits( test_days=self.config.gts.test_days, val_days=self.config.gts.val_days, gap_minutes=self.config.gts.gap_minutes, - compute_ram_fraction=compute_ram_fraction, ) # Mark as completed @@ -235,21 +292,29 @@ def get_splits( return temporal_split, item_index def tensorize(self, split: TemporalSplit) -> TensorizedSplit: - """Tensorize a given TemporalSplit into memory-mapped NumPy arrays. + """Get tensorized splits, computing and caching to disk if not already cached. + + Converts each parquet-backed split into dense `[N, L, K]` memory-mapped NumPy + arrays suitable for out-of-core PyTorch training, where N is the number of + sequences (samples), L is the number of context sessions, and K is the number + of tracks per session. + + Unlike earlier pipeline stages, disk writes are handled internally by + `impact.features.tensorize()`. As memory-mapped arrays are files by definition, + writing to disk and creating the tensor are the same operation. - Converts each Parquet-backed split into dense [N, L, K] arrays suitable for - out-of-core PyTorch training. Results are cached under the pipeline's - `processed_dir`; subsequent calls with a complete cache skip recomputation. + Results are persisted under + `//// + /tensors/` and reused on subsequent calls. Args: split: Temporal dataset splits, as returned by `get_splits()`. Returns: - TensorizedSplit with memory-mapped train, val, and test arrays. + A TensorizedSplit with memory-mapped train, val, and test arrays loaded + from the cached directory. """ - from impact.features import tensorize - - tensors_dir = ( + cache_dir = ( self.processed_dir / self.config.dataset_id / self.config.sessionization_id @@ -257,30 +322,35 @@ def tensorize(self, split: TemporalSplit) -> TensorizedSplit: / self.config.temporal_split_id / "tensors" ) - success_marker = tensors_dir / "_SUCCESS" + success_marker = cache_dir / "_SUCCESS" # Atomicity check if not success_marker.exists(): - logger.info("No complete tensorized splits found in cache, computing...") - # Clean up potentially corrupted partial runs - if tensors_dir.exists(): - shutil.rmtree(tensors_dir) + if cache_dir.exists(): + shutil.rmtree(cache_dir) + cache_dir.mkdir(parents=True) + + logger.info("No complete tensorized splits found in cache, computing...") dataset_by_name = { "train": split.train, "val": split.val, "test": split.test, } - for name, dataset in dataset_by_name.items(): - out_dir = tensors_dir / name - out_dir.mkdir(parents=True) + for name, dataset in tqdm( + dataset_by_name.items(), + desc="Tensorizing", + unit="split", + ): + dataset_cache_dir = cache_dir / name + dataset_cache_dir.mkdir(parents=True) # Discard the result as it has to be recreated at the end either way # (because of caching) _ = tensorize( dataset=dataset, - out_dir=tensors_dir / name, + out_dir=dataset_cache_dir, tracks_per_session=self.config.session.tracks_per_session, context_sessions=self.config.sequence.context_sessions, ) @@ -290,64 +360,7 @@ def tensorize(self, split: TemporalSplit) -> TensorizedSplit: logger.info("Tensorized splits computed successfully") return TensorizedSplit( - train=TensorizedDataset.from_dir(tensors_dir / "train"), - val=TensorizedDataset.from_dir(tensors_dir / "val"), - test=TensorizedDataset.from_dir(tensors_dir / "test"), - ) - - def _save_to_cache( - self, - data: pl.LazyFrame | pl.DataFrame, - cache_dir: Path, - storage_partitions: int = 1, - ) -> None: - """Write dataset to Parquet cache using predefined partitioning strategy. - - The function persists the input dataset to disk either as a single Parquet file - or as multiple partitioned files, depending on the provided `storage_partitions` - parameter. - - If multiple partitions are requested and the dataset contains a `uid` column, - data is split deterministically using a modulo-based partitioning strategy. - Otherwise, the dataset is written as a single Parquet file with a warning. - - Args: - data: Dataset to be cached as Parquet files. - cache_dir: Directory where Parquet files will be saved. - storage_partitions: Number of storage partitions (number of Parquet files). - """ - cache_dir.mkdir(parents=True, exist_ok=True) - - logger.info( - "Saving dataset to cache", - cache_dir=cache_dir.as_posix(), - storage_partitions=storage_partitions, + train=TensorizedDataset.from_dir(cache_dir / "train"), + val=TensorizedDataset.from_dir(cache_dir / "val"), + test=TensorizedDataset.from_dir(cache_dir / "test"), ) - - if storage_partitions == 1: - path = cache_dir / "part_all.parquet" - - logger.info("Writing to a single Parquet file...") - if isinstance(data, pl.LazyFrame): - data.sink_parquet(path) - else: - data.write_parquet(path) - - return - - lf = data.lazy() if isinstance(data, pl.DataFrame) else data - schema = lf.collect_schema() - - if "uid" not in schema.names(): - logger.warning( - "Partitioning by uid not available, trying to write to a single file..." - ) - lf.sink_parquet(cache_dir / "part_all.parquet") - return - - lf = lf.with_columns((pl.col("uid") % storage_partitions).alias("partition")) - for i in range(storage_partitions): - logger.info(f"Saving partition {i + 1}/{storage_partitions}...") - lf.filter(pl.col("partition") == i).collect().write_parquet( - cache_dir / f"part_{i}.parquet" - ) diff --git a/src/impact/pipelines/impact_evaluation_pipeline.py b/src/impact/pipelines/impact_evaluation_pipeline.py index 490a006..68702d8 100644 --- a/src/impact/pipelines/impact_evaluation_pipeline.py +++ b/src/impact/pipelines/impact_evaluation_pipeline.py @@ -27,7 +27,7 @@ import polars as pl from impact.config import ImpactConfig - from impact.schemas import ItemIndex, TensorizedSplit + from impact.schemas import ItemIndex, TensorizedDataset logger = structlog.stdlib.get_logger(__name__) @@ -63,7 +63,7 @@ class _TargetSplits: @final -class IMPACTEvaluationPipeline: +class ImpactEvaluationPipeline: """Late-fusion evaluation pipeline for the IMPACT framework. Combines a trained SASRecNBR model with a deterministic ACT-R scorer via a weighted @@ -86,7 +86,7 @@ def __init__( self, config: ImpactConfig, item_index: ItemIndex, - events_lf: pl.LazyFrame, + listens_lf: pl.LazyFrame, ) -> None: """Initialize the IMPACT Framework evaluation pipeline. @@ -94,20 +94,20 @@ def __init__( config: Complete IMPACT Framework configuration. item_index: Train-derived index mapping raw item (track) IDs to contiguous 1-based indices. - events_lf: LazyFrame of all interaction events; used to build the ACT-R - history index. + listens_lf: LazyFrame of valid listen events, contains the listening history + for all users. Used to build the ACT-R history index. """ self._config = config self._item_index = item_index - self._events_lf = events_lf + self._listens_lf = listens_lf @torch.no_grad() def evaluate( self, - split: TensorizedSplit, + split_partition: TensorizedDataset, sasrec_state_dict_path: Path, ) -> dict[str, float]: - """Evaluate the IMPACT model on the held-out test split. + """Evaluate the IMPACT model on the given split partition. Loads the SASRecNBR checkpoint, builds the ACT-R history index for test users, and runs batched inference over the full item catalog. Final scores are @@ -124,8 +124,9 @@ def evaluate( and `k` is the cutoff. Examples: ndcg_at_5, recall_global_exp_at_10. Args: - split: Tensorized train/val/test temporal splits; only the test partition - is used here. + split_partition: A single tensorized partition to evaluate on. Pass + `TensorizedSplit.val` during hyperparameter search and + `TensorizedSplit.test` for final evaluation on the held-out data. sasrec_state_dict_path: Path to the saved SASRecNBR checkpoint. Returns: @@ -139,8 +140,8 @@ def evaluate( sasrec = self._load_sasrec(sasrec_state_dict_path, device) actr_history_index = ACTRHistoryIndex( - events_lf=self._events_lf, - eval_uids=np.unique(split.test.uid).tolist(), + listens_lf=self._listens_lf, + eval_uids=np.unique(split_partition.uid).tolist(), item_index=self._item_index, ) actr = ACTR( @@ -161,7 +162,7 @@ def evaluate( total_counts: dict[str, int] = dict.fromkeys(get_args(EvalScopeKey), 0) test_loader = DataLoader[EvalSequenceSample]( - dataset=EvalSequenceDataset(split.test), + dataset=EvalSequenceDataset(split_partition), batch_size=self._config.eval.batch_size, shuffle=False, collate_fn=get_eval_collate_fn(), # No negative sampling diff --git a/src/impact/utils.py b/src/impact/utils.py index c74c519..33dd9f0 100644 --- a/src/impact/utils.py +++ b/src/impact/utils.py @@ -1,9 +1,15 @@ from __future__ import annotations +import math +import os +import random from typing import cast +import numpy as np import polars as pl +import psutil import structlog +import torch logger = structlog.stdlib.get_logger(__name__) @@ -16,19 +22,13 @@ def seed_everything(seed: int) -> None: Sets seeds for Python's random, NumPy, and PyTorch (CPU, CUDA, MPS) modules. Also disables cuDNN benchmark mode and enables deterministic operations. Should be called once at the top-level entry point before any data loading or model - initialisation. + initialization. Inspired by: https://gist.github.com/ihoromi4/b681a9088f348942b01711f251e5f964 Args: seed: Global random seed. """ - import os - import random - - import numpy as np - import torch - random.seed(seed) os.environ["PYTHONHASHSEED"] = str(seed) np.random.seed(seed) @@ -47,16 +47,11 @@ def seed_worker(_: int) -> None: # Ignore the argument (worker_id) The per-worker seed is derived from PyTorch's initial worker seed, ensuring workers are both independent and deterministic given the same global seed. - Pass this function as `worker_init_fn` to DataLoader. + Pass this function as `worker_init_fn` to PyTorch's DataLoader. Args: worker_id: Worker index provided automatically by DataLoader. """ - import random - - import numpy as np - import torch - worker_seed = torch.initial_seed() % WORKER_SEED_MODULUS np.random.seed(worker_seed) random.seed(worker_seed) @@ -76,17 +71,12 @@ def estimate_processing_partitions( data: Dataset to be processed. For LazyFrames, size is estimated by sampling the first 1000 rows. target_ram_fraction: Fraction of currently available RAM to target per - partition. It should be keept low (<=0.2) to account for memory spikes + partition. It should be kept low (<=0.2) to account for memory spikes during joins, explodes, and struct operations. Returns: Number of partitions in range [1, 64]. """ - import math - from typing import cast - - import psutil - estimated_size = _estimate_size_bytes(data) available_ram = cast(int, psutil.virtual_memory().available) # type: ignore[redundant-cast] usable_ram = available_ram * target_ram_fraction @@ -114,8 +104,6 @@ def estimate_storage_partitions( `target_file_size_mb` in uncompressed size. Actual files will be smaller due to Parquet compression. - Dataset partitioning can be especially useful for torch dataloaders or numpy memmap. - Args: data: Dataset to be written as partitioned Parquet files. For LazyFrames, size is estimated by sampling the first 1000 rows. @@ -125,8 +113,6 @@ def estimate_storage_partitions( Returns: Number of partitions in range [1, max_partitions]. """ - import math - estimated_size = _estimate_size_bytes(data) target_bytes = target_file_size_mb * 1024**2 partitions = max(1, min(max_partitions, math.ceil(estimated_size / target_bytes)))