| """ |
| Causal step dataset for the Wunder Fund RNN Challenge. |
| |
| Provides CausalStepDataset (left-padded lookback windows for the official |
| stepwise task) and a deterministic train/val loader factory. Targets are |
| aligned so that the example at step t predicts state t+1, and only scored steps |
| (need_prediction=True, with an available next state) are emitted. |
| """ |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset, DataLoader |
| from typing import Optional, List, Tuple |
|
|
| from src.data.protocol import get_feature_columns, load_wunder_dataframe |
| from src.utils.reproducibility import make_torch_generator, seed_worker |
|
|
|
|
| class CausalStepDataset(Dataset): |
| """ |
| Causal online dataset for the official stepwise task. |
| |
| Each sample contains only states up to and including step t and predicts |
| state t+1. Samples are emitted only where need_prediction=True and t+1 is |
| available in the local parquet file. The hidden competition may request a |
| final step prediction whose target is unavailable locally; this dataset |
| deliberately excludes that row from training and local validation. |
| """ |
|
|
| def __init__( |
| self, |
| parquet_path: str, |
| seq_ids: Optional[List[int]] = None, |
| lookback: int = 128, |
| ): |
| self.lookback = int(lookback) |
| if self.lookback <= 0: |
| raise ValueError("lookback must be positive") |
|
|
| df = load_wunder_dataframe(parquet_path, seq_ids=seq_ids) |
| self.feature_cols = get_feature_columns(df) |
| self.n_features = len(self.feature_cols) |
| self.samples = [] |
|
|
| for seq_ix, seq_df in df.groupby("seq_ix", sort=True): |
| states = seq_df[self.feature_cols].to_numpy(dtype=np.float32) |
| need_prediction = seq_df["need_prediction"].to_numpy(dtype=bool) |
| steps = seq_df["step_in_seq"].to_numpy(dtype=np.int64) |
| for pos, step in enumerate(steps): |
| if not need_prediction[pos]: |
| continue |
| target_pos = pos + 1 |
| if target_pos >= len(states): |
| continue |
| self.samples.append( |
| { |
| "seq_ix": int(seq_ix), |
| "step_in_seq": int(step), |
| "states": states, |
| "position": int(pos), |
| "target": states[target_pos], |
| } |
| ) |
|
|
| def __len__(self) -> int: |
| return len(self.samples) |
|
|
| def __getitem__(self, idx: int) -> dict: |
| sample = self.samples[idx] |
| states = sample["states"] |
| pos = sample["position"] |
| start = max(0, pos + 1 - self.lookback) |
| history = states[start : pos + 1] |
|
|
| window = np.zeros((self.lookback, self.n_features), dtype=np.float32) |
| mask = np.zeros((self.lookback,), dtype=np.float32) |
| window[-len(history) :] = history |
| mask[-len(history) :] = 1.0 |
|
|
| return { |
| "seq_ix": torch.tensor(sample["seq_ix"], dtype=torch.long), |
| "step_in_seq": torch.tensor(sample["step_in_seq"], dtype=torch.long), |
| "history": torch.tensor(window, dtype=torch.float32), |
| "history_mask": torch.tensor(mask, dtype=torch.float32), |
| "target": torch.tensor(sample["target"].astype(np.float32, copy=False), dtype=torch.float32), |
| } |
|
|
|
|
| def create_causal_dataloaders( |
| parquet_path: str, |
| train_seq_ids: list[int], |
| val_seq_ids: list[int], |
| *, |
| lookback: int = 128, |
| batch_size: int = 256, |
| num_workers: int = 0, |
| seed: int = 42, |
| ) -> Tuple[DataLoader, DataLoader]: |
| """Create deterministic train/validation loaders for causal step examples.""" |
| train_dataset = CausalStepDataset(parquet_path, seq_ids=train_seq_ids, lookback=lookback) |
| val_dataset = CausalStepDataset(parquet_path, seq_ids=val_seq_ids, lookback=lookback) |
|
|
| train_loader = DataLoader( |
| train_dataset, |
| batch_size=batch_size, |
| shuffle=True, |
| num_workers=num_workers, |
| pin_memory=False, |
| worker_init_fn=seed_worker, |
| generator=make_torch_generator(seed), |
| ) |
| val_loader = DataLoader( |
| val_dataset, |
| batch_size=batch_size, |
| shuffle=False, |
| num_workers=num_workers, |
| pin_memory=False, |
| ) |
| return train_loader, val_loader |
|
|
