File size: 7,807 Bytes
f6044b4 9ed713b f6044b4 3ded97e f6044b4 3ded97e f6044b4 3ded97e f6044b4 9ed713b f6044b4 9ed713b f6044b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """
⚡ ULTRON Systems // Streaming & Sharded Parquet Dataset Loader
Safely streams the 90GB HuggingFace OHLCV dataset month-by-month.
Guarantees O(1) constant disk and RAM usage to prevent Colab quota violations.
"""
import os
import gc
import glob
import shutil
import urllib.request
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import torch
from torch.utils.data import Dataset, DataLoader
from typing import Generator, List, Tuple, Optional
from huggingface_hub import HfFileSystem, hf_hub_download
from config import DatasetConfig
from features import compute_stationary_features
class ShardedOHLCVDataset(Dataset):
"""
In-memory dataset representing windows from a single monthly shard.
Efficiently created from vectorized sliding window views.
"""
def __init__(self, windows: np.ndarray, targets: np.ndarray):
# windows: [N_samples, seq_len, in_features]
# targets: [N_samples, num_horizons]
self.windows = torch.from_numpy(windows).float()
self.targets = torch.from_numpy(targets).float()
def __len__(self):
return len(self.windows)
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
return self.windows[idx], self.targets[idx]
class StreamingShardManager:
"""
Orchestrates downloading, windowing, and memory pruning for monthly Parquet files.
"""
def __init__(self, cfg: DatasetConfig):
self.cfg = cfg
os.makedirs(self.cfg.cache_dir, exist_ok=True)
self.parquet_files = self._discover_remote_shards()
print(f"⚡ [ULTRON] Discovered {len(self.parquet_files)} target monthly shards between {cfg.start_year} and {cfg.end_year}.")
def _discover_remote_shards(self) -> List[str]:
"""Discovers Parquet file paths in HF repo matching year filters."""
try:
fs = HfFileSystem(token=self.cfg.hf_token)
all_files = fs.ls(f"datasets/{self.cfg.repo_id}/data", detail=False)
except Exception as e:
print(f"⚠️ [ULTRON] Direct HfFileSystem query failed ({e}). Constructing standard filenames...")
all_files = []
for y in range(self.cfg.start_year, self.cfg.end_year + 1):
for m in range(1, 13):
all_files.append(f"data/ohlcv_{y}-{m:02d}.parquet")
# Filter matching year range and format
filtered = []
for f in all_files:
fname = os.path.basename(f)
if fname.startswith("ohlcv_") and fname.endswith(".parquet"):
# Extract year
try:
parts = fname.replace("ohlcv_", "").replace(".parquet", "").split("-")
year = int(parts[0])
if self.cfg.start_year <= year <= self.cfg.end_year:
filtered.append(f"data/{fname}")
except Exception:
continue
return sorted(list(set(filtered)))
def stream_monthly_shards(self, start_shard_idx: int = 0) -> Generator[ShardedOHLCVDataset, None, None]:
"""
Yields a ShardedOHLCVDataset for each month, cleaning up disk and RAM immediately after.
Supports resuming from start_shard_idx.
"""
for idx, rel_path in enumerate(self.parquet_files):
if idx < start_shard_idx:
continue
filename = os.path.basename(rel_path)
local_path = os.path.join(self.cfg.cache_dir, filename)
print(f"\n⚡ [ULTRON] Fetching shard: {filename} ...")
try:
# Download single shard (~300MB) directly from HuggingFace with authentication
downloaded_file = hf_hub_download(
repo_id=self.cfg.repo_id,
repo_type="dataset",
filename=rel_path,
local_dir=self.cfg.cache_dir,
token=self.cfg.hf_token,
force_download=False
)
except Exception as e:
print(f"⚠️ [ULTRON] Failed to download {rel_path} via HF Hub: {e}. Trying raw URL fallback...")
url = f"https://huggingface.co/datasets/{self.cfg.repo_id}/resolve/main/{rel_path}"
try:
urllib.request.urlretrieve(url, local_path)
downloaded_file = local_path
except Exception as ex2:
print(f"❌ [ULTRON] Shard {filename} download failed: {ex2}. Skipping.")
continue
# Read and process shard
dataset = self._process_parquet_to_dataset(downloaded_file)
# Remove file from disk immediately to preserve Colab storage quota!
try:
if os.path.exists(downloaded_file):
os.remove(downloaded_file)
except Exception:
pass
if dataset is not None and len(dataset) > 0:
print(f"✅ [ULTRON] Shard {filename} ingested: Generated {len(dataset)} overlapping windows.")
yield dataset
else:
print(f"⚠️ [ULTRON] Shard {filename} yielded no valid samples. Proceeding to next shard.")
# Trigger garbage collection
gc.collect()
def _process_parquet_to_dataset(self, file_path: str) -> Optional[ShardedOHLCVDataset]:
"""Reads parquet, computes stationary features, and extracts sliding windows."""
try:
# Read only required columns
table = pq.read_table(
file_path,
columns=["timestamp", "open", "high", "low", "close", "volume", "ticker"]
)
df = table.to_pandas()
del table
gc.collect()
except Exception as e:
print(f"❌ Error reading parquet file {file_path}: {e}")
return None
# Filter tickers if specified
if self.cfg.liquid_tickers:
df = df[df["ticker"].isin(self.cfg.liquid_tickers)]
if len(df) == 0:
del df
return None
all_windows = []
all_targets = []
# Group by ticker to prevent cross-ticker leakage
for ticker, ticker_df in df.groupby("ticker"):
feats, targets = compute_stationary_features(ticker_df, self.cfg.pred_horizons)
n_rows = len(feats)
if n_rows < self.cfg.seq_len:
continue
# Vectorized overlapping sliding window slicing:
# Stride over sequence with window_stride (e.g. 16 minutes)
num_windows = (n_rows - self.cfg.seq_len) // self.cfg.window_stride + 1
if num_windows <= 0:
continue
for i in range(0, n_rows - self.cfg.seq_len + 1, self.cfg.window_stride):
w_feats = feats[i : i + self.cfg.seq_len]
# Target corresponds to the future return from the final bar of this window
w_target = targets[i + self.cfg.seq_len - 1]
# Check for NaNs or Inf
if not np.isnan(w_target).any() and not np.isnan(w_feats).any():
all_windows.append(w_feats)
all_targets.append(w_target)
del df
gc.collect()
if not all_windows:
return None
windows_arr = np.array(all_windows, dtype=np.float32)
targets_arr = np.array(all_targets, dtype=np.float32)
return ShardedOHLCVDataset(windows_arr, targets_arr)
|