import modal import os import uuid from datetime import datetime import yfinance as yf import pandas as pd import requests import io import os import math import json import warnings import random import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from datetime import datetime from typing import Optional, Tuple, Dict, List, Union from dataclasses import dataclass, asdict from pathlib import Path from torch.utils.data import Dataset, DataLoader from torch.amp import autocast, GradScaler from transformers import AutoModelForSeq2SeqLM, AutoConfig from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score import matplotlib.pyplot as plt # --- Configuration --- DOWNLOAD_DATA = False RAW_DATA_DIR = "raw_data" os.makedirs(RAW_DATA_DIR, exist_ok=True) # Global feature order for training consistency # LCS Upgrade (Phase 2): All covariates are now scale-invariant ratios COV_COLS = [ # Price ratios: (price / SMA) - 1 → mean-reverting, scale-invariant "ratio_mva5", "ratio_mva10", "ratio_mva15", "ratio_mva20", "ratio_mva30", # Volume normalized by SMA-20 (added inside add_technical_indicators) "Volume_ratio", # RSI is already [0,100] and scale-invariant "rsi14", # Bollinger %B: (Close - Lower) / (Upper - Lower) → [0,1] range, scale-invariant "BB_PercentB", # MACD normalised by Close price → scale-invariant "MACD_norm", "Signal_norm", "Hist_norm", ] START_DATE = "2013-01-01" HEADERS = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36" } def run_download_pipeline(): """Encapsulated downloader logic.""" if not DOWNLOAD_DATA: print(f"Skipping download (DOWNLOAD_DATA={DOWNLOAD_DATA}).") return # 1. Get Symbols symbols = get_combined_market_tickers() if len(symbols) > 0: # 2. Download Data historical_data = get_stock_data(symbols, start_date=START_DATE) # 3. Save Files print(f"\nSaving {len(historical_data)} files to {RAW_DATA_DIR}...") for ticker, df in historical_data.items(): file_path = os.path.join(RAW_DATA_DIR, f"{ticker}.csv") df.to_csv(file_path) print("Done.") else: print("Failed to retrieve symbol list.") def get_wiki_tickers(url, table_keywords): """ Helper to fetch ticker symbols from a Wikipedia URL table. """ tickers = [] try: response = requests.get(url, headers=HEADERS) response.raise_for_status() tables = pd.read_html(io.StringIO(response.text)) target_table = None for table in tables: cols = [c.lower() for c in table.columns] if any(k in cols for k in ['symbol', 'ticker']): target_table = table break if target_table is not None: # Find the specific column name (e.g. "Symbol" or "Ticker") col_name = next(c for c in target_table.columns if c.lower() in ['symbol', 'ticker']) tickers = target_table[col_name].tolist() except Exception as e: print(f" -> Error scraping {url}: {e}") return tickers def get_combined_market_tickers(): """ Aggregates S&P 500, NASDAQ 100, and S&P 400 (MidCap). """ print("--- Generaring Symbol List ---") all_tickers = set() # 1. S&P 500 sp500 = get_wiki_tickers("https://en.wikipedia.org/wiki/List_of_S%26P_500_companies", ['symbol']) all_tickers.update(sp500) print(f"S&P 500 count: {len(sp500)}") # 2. NASDAQ 100 ndx100 = get_wiki_tickers("https://en.wikipedia.org/wiki/Nasdaq-100", ['ticker', 'symbol']) all_tickers.update(ndx100) print(f"NASDAQ 100 count: {len(ndx100)}") # 3. S&P 400 MidCap (Ensures >250 quality stocks) sp400 = get_wiki_tickers("https://en.wikipedia.org/wiki/List_of_S%26P_400_companies", ['ticker', 'symbol']) all_tickers.update(sp400) print(f"S&P 400 count: {len(sp400)}") # Clean: Replace '.' with '-' (e.g. BRK.B -> BRK-B) and remove non-strings clean_list = [str(t).replace('.', '-') for t in all_tickers if isinstance(t, str)] unique_list = sorted(list(set(clean_list))) print(f"\nTotal unique symbols found: {len(unique_list)}") return unique_list def get_stock_data(tickers, start_date=None, interval="1d", end_date=None): """ Downloads data. If start_date is provided, it uses that. Otherwise, it uses period="max". """ if not tickers: print("No tickers provided.") return {} if end_date is None: end_date = datetime.now().strftime('%Y-%m-%d') print(f"Starting download for {len(tickers)} symbols...") if start_date: print(f"Data Range: {start_date} to {end_date}") else: print(f"Data Range: MAX available history to {end_date}") all_data = {} BATCH_SIZE = 50 for i in range(0, len(tickers), BATCH_SIZE): batch = tickers[i:i + BATCH_SIZE] print(f"Processing batch {i} - {i+len(batch)}...") try: # --- NEW LOGIC FOR START DATE --- # We construct the arguments dynamically. download_args = { "tickers": ' '.join(batch), "interval": interval, "end": end_date, "group_by": 'ticker', "auto_adjust": True, "progress": False, "threads": True } # If a start date is given, use 'start', otherwise use 'period="max"' if start_date: download_args["start"] = start_date else: download_args["period"] = "max" data = yf.download(**download_args) if not data.empty: if len(batch) == 1: all_data[batch[0]] = data.dropna(how='all') else: for ticker in batch: try: t_data = data.get(ticker) if t_data is not None and not t_data.empty: # Filter out rows that are all NaN t_data = t_data.dropna(how='all') if len(t_data) > 0: all_data[ticker] = t_data except Exception: continue except Exception as e: print(f"Batch error: {e}") return all_data TRAIN_END = pd.Timestamp("2021-12-31") VAL_END = pd.Timestamp("2023-12-31") # ============================================================================= # SECTION A: TECHNICAL INDICATORS # ============================================================================= def simple_moving_average( df: pd.DataFrame, period: int, price_col: str = "Close", return_list: bool = False ): """Simple Moving Average (SMA) over a right-aligned rolling window.""" if price_col not in df.columns or not isinstance(period, int) or period <= 0: raise ValueError("Invalid inputs.") sma = df[price_col].rolling(window=period, min_periods=period).mean() return sma.tolist() if return_list else sma def ema( df: pd.DataFrame, period: int, price_col: str = "Close", seed: str = "price", # kept for API compatibility min_periods: Optional[int] = None, return_list: bool = False ): """Exponential Moving Average (EMA).""" if price_col not in df.columns or not isinstance(period, int) or period <= 0: raise ValueError("Invalid inputs.") x = df[price_col].astype(float) # Standard EMA (recursive, no adjust) ser = x.ewm(span=period, adjust=False, min_periods=min_periods).mean() return ser.tolist() if return_list else ser def rsi_wilder( df: pd.DataFrame, period: int = 14, price_col: str = "Close", return_list: bool = False ): """Wilder's Relative Strength Index (RSI).""" if price_col not in df.columns or not isinstance(period, int) or period <= 1: raise ValueError("Invalid inputs.") price = df[price_col].astype(float) delta = price.diff() gains = delta.clip(lower=0) losses = (-delta).clip(lower=0) avg_gain = gains.ewm(alpha=1/period, adjust=False, min_periods=period).mean() avg_loss = losses.ewm(alpha=1/period, adjust=False, min_periods=period).mean() rs = avg_gain / avg_loss rsi = 100.0 - (100.0 / (1.0 + rs)) rsi = rsi.where(~((avg_loss == 0) & (avg_gain > 0)), 100.0) rsi = rsi.where(~((avg_gain == 0) & (avg_loss > 0)), 0.0) rsi = rsi.where(~((avg_gain == 0) & (avg_loss == 0)), 50.0) return rsi.tolist() if return_list else rsi def bollinger_bands( df: pd.DataFrame, period: int = 20, num_std_dev: float = 2.0, price_col: str = "Close", ddof: int = 0, return_extras: bool = False ) -> pd.DataFrame: """Bollinger Bands over a right-aligned rolling window.""" if price_col not in df.columns or not isinstance(period, int) or period <= 1: raise ValueError("Invalid inputs.") s = df[price_col].astype(float) middle = s.rolling(window=period, min_periods=period).mean() sigma = s.rolling(window=period, min_periods=period).std(ddof=ddof) upper = middle + num_std_dev * sigma lower = middle - num_std_dev * sigma out = pd.DataFrame({ "BB_Middle": middle, "BB_Upper": upper, "BB_Lower": lower }, index=df.index) if return_extras: width = upper - lower percent_b = (s - lower) / width percent_b = percent_b.where(width > 0) bandwidth = width / middle out["BB_PercentB"] = percent_b out["BB_Bandwidth"] = bandwidth return out def macd( df: pd.DataFrame, price_col: str = "Close", fast_period: int = 12, slow_period: int = 26, signal_period: int = 9, min_periods: Optional[int] = None ) -> pd.DataFrame: """MACD (Moving Average Convergence Divergence).""" if price_col not in df.columns or not (fast_period < slow_period and slow_period > 0): raise ValueError("Invalid periods/columns.") x = pd.to_numeric(df[price_col], errors="coerce").astype(float) ema_fast = x.ewm(span=fast_period, adjust=False).mean() ema_slow = x.ewm(span=slow_period, adjust=False).mean() macd_line = ema_fast - ema_slow signal = macd_line.ewm(span=signal_period, adjust=False).mean() histogram = macd_line - signal return pd.DataFrame( {"MACD": macd_line, "Signal_Line": signal, "MACD_Histogram": histogram}, index=df.index ) def pivot_points_fibonacci_FIXED( df: pd.DataFrame, high_col: str = "High", low_col: str = "Low", close_col: str = "Close", session: str = "D" ) -> pd.DataFrame: """ Fibonacci Pivot Points, shifted one period forward to avoid leakage. """ df = df.sort_index() ohlc = df[[high_col, low_col, close_col]].copy() sess = ohlc.resample(session, label="left", closed="left").agg({ high_col: "max", low_col: "min", close_col: "last" }) H, L, C = sess[high_col], sess[low_col], sess[close_col] P = (H + L + C) / 3.0 R = (H - L) R1 = P + 0.382 * R R2 = P + 0.618 * R R3 = P + 1.000 * R S1 = P - 0.382 * R S2 = P - 0.618 * R S3 = P - 1.000 * R fib_sess = pd.DataFrame({ "PP": P, "R1": R1, "R2": R2, "R3": R3, "S1": S1, "S2": S2, "S3": S3, }, index=sess.index) fib_sess_shifted = fib_sess.shift(1) return fib_sess_shifted.reindex(df.index).ffill() def add_technical_indicators(df: pd.DataFrame) -> pd.DataFrame: out = df.copy() # Raw SMAs (needed for ratio computation, not stored as raw price cols) sma5 = simple_moving_average(out, 5) sma10 = simple_moving_average(out, 10) sma15 = simple_moving_average(out, 15) sma20 = simple_moving_average(out, 20) sma30 = simple_moving_average(out, 30) # LCS Upgrade (Phase 2): Price/SMA ratios → scale-invariant out["ratio_mva5"] = (out["Close"] / sma5) - 1 out["ratio_mva10"] = (out["Close"] / sma10) - 1 out["ratio_mva15"] = (out["Close"] / sma15) - 1 out["ratio_mva20"] = (out["Close"] / sma20) - 1 out["ratio_mva30"] = (out["Close"] / sma30) - 1 # Volume normalized by 20-day SMA of Volume (scale-invariant) if "Volume" in out.columns: vol_sma20 = out["Volume"].rolling(window=20, min_periods=20).mean() out["Volume_ratio"] = out["Volume"] / (vol_sma20 + 1e-8) else: out["Volume_ratio"] = 0.0 # fallback if Volume absent # RSI: already [0, 100], scale-invariant out["rsi14"] = rsi_wilder(out, 14) # Bollinger Bands: use %B instead of raw upper/lower bb = bollinger_bands(out, period=20, num_std_dev=2.0, price_col="Close", return_extras=True) out = out.join(bb) # BB_PercentB already computed by return_extras=True; handle NaN (flat price sections) out["BB_PercentB"] = out["BB_PercentB"].fillna(0.5) # MACD normalised by Close price → scale-invariant macd_df = macd(out, price_col="Close") out["MACD_norm"] = macd_df["MACD"] / (out["Close"] + 1e-8) out["Signal_norm"] = macd_df["Signal_Line"] / (out["Close"] + 1e-8) out["Hist_norm"] = macd_df["MACD_Histogram"]/ (out["Close"] + 1e-8) return out # ============================================================================= # SECTION B: WINDOWING AND MULTI-TICKER LOADER # ============================================================================= def build_windows_from_df(df: pd.DataFrame, T: int = 64, H: int = 16, columns: list = None): """ LCS Upgrade (Phase 3): Produces sliding windows with Local Context Scaling. For each window: 1. Extract raw price history (T steps) + future (H steps). 2. Compute local scale = mean(abs(history prices)). 3. Scale the entire window by this local scale → scaled prices ≈ N(1, σ). 4. Tokenize scaled_history as input_ids, scaled_future as labels. 5. Return raw prices, covariates, raw labels, AND local scales. Windows where scale <= 0.001 (flat/zero data) are dropped. """ prices = df["price"].values.astype("float32") # Use passed columns (from checkpoint) or global fallback target_cols = columns if columns is not None else COV_COLS missing = [c for c in target_cols if c not in df.columns] if missing: raise ValueError(f"Missing columns in dataframe: {missing}") cov = df[target_cols].values.astype("float32") N_total = len(df) windows = N_total - T - H + 1 if windows <= 0: raise ValueError("Dataset too small for given T and H.") price_windows = [] cov_windows = [] label_windows = [] scale_windows = [] for i in range(windows): p_window = prices[i:i+T] f_window = prices[i+T:i+T+H] # Filter non-positive windows early if np.any(p_window <= 0) or np.any(f_window <= 0): continue # LCS: compute local scale from history local_scale = float(np.mean(np.abs(p_window))) if local_scale <= 0.001: continue # Drop flat/near-zero windows price_windows.append(p_window) cov_windows.append(cov[i:i+T]) label_windows.append(f_window) scale_windows.append(local_scale) if not price_windows: return np.array([]), np.array([]), np.array([]), np.array([]) return ( np.stack(price_windows), np.stack(cov_windows), np.stack(label_windows), np.array(scale_windows, dtype="float32"), ) def load_multi_csvs_stream( data_dir: str, T: int = 64, H: int = 16, min_len: int = 120, max_files: int = 200 # <-- limit processed symbols per session ): """ Stream CSVs in small batches to avoid RAM crashes. Returns concatenated arrays but processes at most `max_files` tickers. """ all_prices = [] all_cov = [] all_labels = [] files = [f for f in os.listdir(data_dir) if f.endswith(".csv")] files = files[:max_files] # HARD LIMIT print(f"Processing {len(files)} tickers instead of full dataset...") for fname in files: path = os.path.join(data_dir, fname) try: df = pd.read_csv(path, parse_dates=["Date"]) except Exception: continue if "Date" not in df.columns: continue df = df.sort_values("Date") df = df.drop_duplicates() df = df.set_index("Date") # Required cols if not {"Open","High","Low","Close"}.issubset(df.columns): continue # Business day alignment #df = df.resample("B").asfreq().ffill().bfill() df = df.resample("B").asfreq().ffill() df = df.dropna() df["price"] = df["Close"] if len(df) < min_len: continue # Compute indicators (heavy! but OK per file) df = add_technical_indicators(df) df = df.dropna() if len(df) < T + H: continue # Build windows — now returns (prices, cov, labels, scales) result = build_windows_from_df(df, T=T, H=H) p_i, c_i, y_i = result[0], result[1], result[2] # Only append if windows were generated if len(p_i) > 0: all_prices.append(p_i) all_cov.append(c_i) all_labels.append(y_i) if not all_prices: raise ValueError("No data after streaming.") prices = np.concatenate(all_prices, axis=0) covs = np.concatenate(all_cov, axis=0) labels = np.concatenate(all_labels, axis=0) return prices, covs, labels def compute_mase_scale_multi(data_dir: str) -> float: """ Compute naive random-walk MAE across all tickers in data_dir for MASE scaling. """ required_cols = ["Open", "High", "Low", "Close"] diffs = [] for fname in os.listdir(data_dir): if not fname.endswith(".csv"): continue path = os.path.join(data_dir, fname) try: df = pd.read_csv(path, parse_dates=["Date"]) except Exception: continue if "Date" not in df.columns: continue df = df.sort_values("Date") df = df.drop_duplicates() df = df.set_index("Date") if not set(required_cols).issubset(df.columns): continue df = df[required_cols].copy() #df = df.resample("B").asfreq().ffill().bfill() df = df.resample("B").asfreq().ffill() df = df.dropna() df["price"] = df["Close"] p = df["price"].values.astype("float32") if len(p) < 2: continue diffs.append(np.abs(p[1:] - p[:-1])) if not diffs: return 1.0 diffs_all = np.concatenate(diffs, axis=0) mase_scale = float(np.mean(diffs_all)) return max(mase_scale, 1e-8) # ============================================================================= # SECTION C: TOKENIZER & NORMALIZER # ============================================================================= class SimpleChronosTokenizer: """ LCS Upgrade: Local Context Scaling tokenizer. Operates on scaled prices (price / local_scale), covering [0.0, 5.0]. Replaces the old log-return tokenizer to eliminate compounding errors. """ def __init__(self, num_bins: int = 1024, pad_token_id: int = 0, min_val: float = 0.0, max_val: float = 5.0): self.num_bins = num_bins self.pad_token_id = pad_token_id self.min_val = min_val self.max_val = max_val self.range = self.max_val - self.min_val + 1e-8 self.bos_token_id = 1 self.eos_token_id = 2 self.vocab_offset = 3 self.vocab_size = num_bins + self.vocab_offset self.fitted = True def fit(self, prices: np.ndarray): # LCS: No-op. Fixed range [0.0, 5.0] covers scaled prices well. print(f"✓ Tokenizer using fixed LCS range [{self.min_val}, {self.max_val}]") def encode(self, arr: np.ndarray) -> np.ndarray: arr = np.array(arr, dtype=float) # Clamp; warn if many values are out-of-range (sign of bad scale) clamp_count = np.sum(arr > self.max_val) if clamp_count > len(arr) * 0.1: warnings.warn(f"⚠️ {clamp_count}/{len(arr)} values clamped above {self.max_val}. " "Consider re-checking local scale computation.") arr = np.clip(arr, self.min_val, self.max_val) scaled = (arr - self.min_val) / self.range scaled = np.clip(scaled, 0, 1) bins = (scaled * (self.num_bins - 1)).astype(np.int64) tokens = bins + self.vocab_offset tokens[~np.isfinite(arr)] = self.pad_token_id return tokens def decode(self, tokens: np.ndarray) -> np.ndarray: tokens = np.array(tokens, dtype=np.int64) mask_special = ( (tokens == self.pad_token_id) | (tokens == self.bos_token_id) | (tokens == self.eos_token_id) ) bins = np.clip(tokens - self.vocab_offset, 0, self.num_bins - 1) scaled = bins.astype(float) / (self.num_bins - 1) values = scaled * self.range + self.min_val values[mask_special] = np.nan return values def save_config(self, path: str): config = { "num_bins": self.num_bins, "pad_token_id": self.pad_token_id, "min_val": self.min_val, "max_val": self.max_val, "mode": "local_scale" # LCS Upgrade: renamed from "log_return" } with open(path, "w") as f: json.dump(config, f, indent=2) print(f"✓ Tokenizer config saved to {path}") @classmethod def load_config(cls, path: str): with open(path, "r") as f: config = json.load(f) return cls( num_bins=config["num_bins"], pad_token_id=config["pad_token_id"], min_val=config["min_val"], max_val=config["max_val"] ) class ZScoreNormalizer: """ Anti-leakage Z-score normalizer with persistence. """ def __init__(self, epsilon: float = 1e-8): self.mean: Optional[torch.Tensor] = None self.std: Optional[torch.Tensor] = None self.num_features: Optional[int] = None self.epsilon = epsilon self.fitted = False def fit(self, train_data: torch.Tensor): if train_data.dim() != 3: raise ValueError(f"Expected (N, T, F), got {train_data.shape}") self.mean = train_data.mean(dim=(0, 1), keepdim=True) self.std = train_data.std(dim=(0, 1), keepdim=True) + self.epsilon self.num_features = train_data.shape[2] self.fitted = True low_var = (self.std.squeeze() < (self.epsilon * 10)).cpu() if low_var.any(): idx = torch.where(low_var)[0].tolist() warnings.warn(f"⚠️ Low variance features at indices: {idx}") print(f"✓ Normalizer fitted on training covariates (F={self.num_features})") def transform(self, data: torch.Tensor) -> torch.Tensor: if not self.fitted: raise ValueError("Normalizer not fitted. Call fit() first.") if data.shape[-1] != self.num_features: raise ValueError(f"Feature mismatch: expected {self.num_features}, got {data.shape[-1]}") return (data - self.mean.to(data.device)) / self.std.to(data.device) def inverse_transform(self, data: torch.Tensor) -> torch.Tensor: if not self.fitted: raise ValueError("Normalizer not fitted.") return data * self.std.to(data.device) + self.mean.to(data.device) # --- NEW PERSISTENCE METHODS --- def save(self, dir_path: str): if not self.fitted: raise ValueError("Cannot save unfitted normalizer") # Save tensors torch.save(self.mean, os.path.join(dir_path, "norm_mean.pt")) torch.save(self.std, os.path.join(dir_path, "norm_std.pt")) # Save metadata meta = {"num_features": self.num_features, "epsilon": self.epsilon} with open(os.path.join(dir_path, "norm_meta.json"), "w") as f: json.dump(meta, f, indent=2) print(f"✓ Normalizer stats saved to {dir_path}") @classmethod def load(cls, dir_path: str): meta_path = os.path.join(dir_path, "norm_meta.json") if not os.path.exists(meta_path): raise FileNotFoundError(f"Missing normalizer metadata: {meta_path}") with open(meta_path, "r") as f: meta = json.load(f) obj = cls(epsilon=meta["epsilon"]) obj.num_features = meta["num_features"] # Load tensors (map to CPU initially) obj.mean = torch.load(os.path.join(dir_path, "norm_mean.pt"), map_location="cpu") obj.std = torch.load(os.path.join(dir_path, "norm_std.pt"), map_location="cpu") obj.fitted = True print(f"✓ Normalizer loaded (F={obj.num_features})") return obj # ============================================================================= # SECTION D: DUAL-PATH CROSS-ATTENTION ADAPTER (PHASE 2) # ============================================================================= import torch import torch.nn as nn import torch.nn.functional as F import math class xLSTMCovariateEncoder(nn.Module): """ Strategy 1: Deep xLSTM Encoder — 3 stacked Pre-Norm residual mLSTM blocks. Layer 1 (Perception): raw indicators -> baseline market states Layer 2 (Context): combine states across time Layer 3 (Translation): map to T5 cross-attention space Each layer: residual highway + branch(LayerNorm -> mLSTMBlock) """ def __init__(self, input_dim=11, d_model=256, num_heads=8, num_layers=3, expansion_factor=2, dropout=0.1): super().__init__() self.input_dim = input_dim self.d_model = d_model self.num_heads = num_heads self.num_layers = num_layers # Input projection: input_dim -> d_model self.projection = nn.Linear(input_dim, d_model) # Sinusoidal positional encoding self.pos_encoding = SinusoidalPositionalEncoding(d_model) # Stack of Pre-Norm mLSTM blocks self.norms = nn.ModuleList([nn.LayerNorm(d_model) for _ in range(num_layers)]) self.blocks = nn.ModuleList([ MultiHeadmLSTMBlock(d_model=d_model, num_heads=num_heads, expansion_factor=expansion_factor) for _ in range(num_layers) ]) self.dropout = nn.Dropout(dropout) def forward(self, covariates): """ Args: covariates: (B, T, input_dim) Returns: output: (B, T, d_model) """ # 1. Project + positional encoding x = self.projection(covariates) # (B, T, d_model) x = self.pos_encoding(x) # (B, T, d_model) # 2. Hierarchical Pre-Norm residual stack for norm, block in zip(self.norms, self.blocks): # Branch: LayerNorm -> mLSTM -> Dropout branch = self.dropout(block(norm(x))) # Residual highway (untouched) x = x + branch return x # (B, T, d_model) class MultiHeadmLSTMBlock(nn.Module): """ Multi-head mLSTM block with Up-Projection + SiLU gate (Strategy 1). Operates in expanded space (d_model * expansion_factor) internally, then projects back down to d_model. """ def __init__(self, d_model=256, num_heads=8, expansion_factor=2): super().__init__() assert d_model % num_heads == 0 self.d_model = d_model self.num_heads = num_heads self.d_expanded = d_model * expansion_factor # 512 default assert self.d_expanded % num_heads == 0 self.d_head = self.d_expanded // num_heads # 64 per head # Up-projection: d_model -> d_expanded (Cover's Theorem) self.up_proj = nn.Linear(d_model, self.d_expanded) # Per-head projections operating in EXPANDED space self.W_q = nn.Parameter(torch.randn(num_heads, self.d_head, self.d_head)) self.W_k = nn.Parameter(torch.randn(num_heads, self.d_head, self.d_head)) self.W_v = nn.Parameter(torch.randn(num_heads, self.d_head, self.d_head)) self.W_alpha = nn.Parameter(torch.randn(num_heads, self.d_head, 1)) self.b_alpha = nn.Parameter(torch.zeros(num_heads, 1)) self.W_out = nn.Parameter(torch.randn(num_heads, self.d_head, self.d_head)) # Down-projection: d_expanded -> d_model self.down_proj = nn.Linear(self.d_expanded, d_model) # Post-memory norm (operates on d_head in expanded space) self.post_norm = nn.LayerNorm(self.d_head) self._init_weights() def _init_weights(self): nn.init.xavier_uniform_(self.W_q) nn.init.xavier_uniform_(self.W_k) nn.init.xavier_uniform_(self.W_v) nn.init.xavier_uniform_(self.W_alpha) nn.init.xavier_uniform_(self.W_out) nn.init.uniform_(self.b_alpha, 0.0, 0.5) nn.init.xavier_uniform_(self.up_proj.weight) nn.init.xavier_uniform_(self.down_proj.weight) def forward(self, x): """ Args: x: (B, T, d_model) Returns: (B, T, d_model) """ B, T, _ = x.shape # 1. Up-project + SiLU filter (information gate before memory write) x_up = F.silu(self.up_proj(x)) # (B, T, d_expanded) # 2. Split into heads in expanded space x_heads = x_up.view(B, T, self.num_heads, self.d_head).transpose(1, 2) # (B, H, T, d_head) # 3. Q, K, V projections Q = torch.einsum('bhtd,hde->bhte', x_heads, self.W_q) K = torch.einsum('bhtd,hde->bhte', x_heads, self.W_k) V = torch.einsum('bhtd,hde->bhte', x_heads, self.W_v) # 4. Gating alpha_logits = torch.einsum('bhtd,hdf->bhtf', x_heads, self.W_alpha) alpha_logits = alpha_logits + self.b_alpha.view(1, self.num_heads, 1, 1) alpha = torch.clamp(torch.sigmoid(alpha_logits), min=1e-6, max=1.0 - 1e-6).unsqueeze(-1) # 5. Rank-1 memory updates updates = torch.matmul(K.unsqueeze(-1), V.unsqueeze(-2)) # (B, H, T, d, d) # 6. Sequential memory update M = self._sequential_memory_update(alpha.transpose(1, 2), updates.transpose(1, 2)) M = M.transpose(1, 2) # (B, H, T, d, d) # 7. Memory read + post-norm h = torch.matmul(M, Q.unsqueeze(-1)).squeeze(-1) # (B, H, T, d_head) h = self.post_norm(h) # 8. Output projection with SiLU y = F.silu(torch.einsum('bhtd,hde->bhte', h, self.W_out)) # (B, H, T, d_head) # 9. Reassemble heads -> d_expanded y = y.transpose(1, 2).contiguous().view(B, T, self.d_expanded) # 10. Down-project back to d_model return self.down_proj(y) # (B, T, d_model) def _sequential_memory_update(self, alpha, updates): B, T, H, d, _ = updates.shape memory_states = [] m_curr = alpha[:, 0] * updates[:, 0] memory_states.append(m_curr) for t in range(1, T): m_prev = memory_states[-1] m_curr = (1 - alpha[:, t]) * m_prev + alpha[:, t] * updates[:, t] memory_states.append(m_curr) return torch.stack(memory_states, dim=1) class SinusoidalPositionalEncoding(nn.Module): """Standard sinusoidal positional encoding.""" def __init__(self, d_model, max_len=5000): super().__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) # (1, max_len, d_model) self.register_buffer('pe', pe) def forward(self, x): """ Args: x: (B, T, d_model) Returns: x + positional encoding: (B, T, d_model) """ return x + self.pe[:, :x.size(1), :] class SinusoidalTimeEncoding(nn.Module): """ Robust Sinusoidal PE that handles odd/even dimensions correctly. """ def __init__(self, d_model: int, max_len: int = 500): super().__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) # Handle odd d_model by slicing the cosine term n_cos = pe[:, 1::2].shape[1] pe[:, 1::2] = torch.cos(position * div_term)[:, :n_cos] self.register_buffer('pe', pe.unsqueeze(0)) def forward(self, x: torch.Tensor) -> torch.Tensor: T = x.size(1) return x + self.pe[:, :T, :] class CovariateEncoder(nn.Module): def __init__(self, input_dim, d_model, nhead=4, num_layers=1, dropout=0.1): """ Encodes raw covariates into a dense, time-aware representation. Input: (B, T, F) -> Output: (B, T, d_model) """ super().__init__() # 1. Project features to model dimension self.input_proj = nn.Linear(input_dim, d_model) # 2. Add Time Awareness (Crucial for Transformer) self.pos_encoder = SinusoidalTimeEncoding(d_model) # 3. Transformer Encoder to mix temporal features encoder_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=d_model * 2, dropout=dropout, batch_first=True ) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) def forward(self, covariates): # covariates: [B, T, F] x = self.input_proj(covariates) # [B, T, D] x = self.pos_encoder(x) # Add PE return self.transformer(x) # [B, T, D] class CrossAttentionAdapter(nn.Module): def __init__(self, d_model, nhead=4, dropout=0.1): """ Injects covariate context into encoder hidden states via Cross-Attention. Query = Encoder Hidden States Key/Value = Covariate Context """ super().__init__() self.cross_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True, dropout=dropout) self.layernorm = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) # 🟢 Zero-Init Output Projection (Optional but recommended for stability) # This helps the model start close to Identity behavior nn.init.zeros_(self.cross_attn.out_proj.weight) nn.init.zeros_(self.cross_attn.out_proj.bias) def forward(self, encoder_hidden, cov_context): # encoder_hidden: [B, T, D] (Query) # cov_context: [B, T, D] (Key/Value) residual = encoder_hidden # Cross Attention: Q=Hidden, K=Cov, V=Cov attn_output, _ = self.cross_attn( query=encoder_hidden, key=cov_context, value=cov_context, need_weights=False ) # Residual Connection + Norm out = self.layernorm(residual + self.dropout(attn_output)) return out class DualPathAdapter(nn.Module): def __init__(self, cov_dim, d_model, num_layers=3, expansion_factor=2): super().__init__() self.cov_encoder = xLSTMCovariateEncoder( input_dim=cov_dim, d_model=d_model, num_heads=8, num_layers=num_layers, # <--- Dynamic variable (Only once) expansion_factor=expansion_factor, # <--- Dynamic variable (Only once) dropout=0.1 ) # No cov_to_t5_projection needed: encoder output is already d_model self.cross_adapter = CrossAttentionAdapter(d_model=d_model) def forward(self, covariates, encoder_hidden): cov_encoded = self.cov_encoder(covariates) # (B, T, d_model) directly return self.cross_adapter(encoder_hidden, cov_encoded) class ChronosDualPath(nn.Module): def __init__( self, t5_model: AutoModelForSeq2SeqLM, num_features: int, freeze_encoder: bool = True, num_layers: int = 3, expansion_factor: int = 2 ): super().__init__() self.t5_model = t5_model self.d_model = t5_model.config.d_model # Initialize the Dual-Path Adapter self.adapter = DualPathAdapter( cov_dim=num_features, d_model=self.d_model, num_layers=num_layers, expansion_factor=expansion_factor ) # Apply freezing configuration immediately self.configure_t5_freezing(freeze_encoder) def configure_t5_freezing(self, freeze_encoder: bool): """ Option A (freeze_encoder=True): Encoder Frozen, Decoder Unfrozen. Option B (freeze_encoder=False): ALL Unfrozen (Full Fine-Tuning). """ # 1. Handle Encoder for p in self.t5_model.encoder.parameters(): p.requires_grad = not freeze_encoder # 2. Handle Decoder (Always unfrozen to learn new target distribution) for p in self.t5_model.decoder.parameters(): p.requires_grad = True if freeze_encoder: print("✓ T5 Freezing: Option A (Encoder FROZEN, Decoder UNFROZEN)") else: print("✓ T5 Freezing: Option B (Full Fine-Tuning, ALL UNFROZEN)") #change def forward( self, input_ids: torch.Tensor, price: torch.Tensor, cov: torch.Tensor, labels: Optional[torch.Tensor] = None, decoder_input_ids: Optional[torch.Tensor] = None, # Required for Step E1 attention_mask: Optional[torch.Tensor] = None, decoder_attention_mask: Optional[torch.Tensor] = None, # Future-proofing ): encoder_outputs = self.t5_model.encoder( input_ids=input_ids, attention_mask=attention_mask, return_dict=True ) encoder_outputs.last_hidden_state = self.adapter( covariates=cov, encoder_hidden=encoder_outputs.last_hidden_state ) # Pass explicit decoder inputs and masks outputs = self.t5_model( encoder_outputs=encoder_outputs, labels=labels, decoder_input_ids=decoder_input_ids, attention_mask=attention_mask, decoder_attention_mask=decoder_attention_mask, return_dict=True ) return outputs # --- Generation Utilities --- def _restrict_logits_to_numeric(self, logits): start = self.t5_model.config.vocab_offset end = start + self.t5_model.config.num_bins mask = torch.full_like(logits, float("-inf")) mask[..., start:end] = logits[..., start:end] return mask def _sample_numeric(self, logits, temperature=1.0, top_p=0.9): logits = logits / max(temperature, 1e-8) probs = torch.softmax(logits, dim=-1) sorted_probs, sorted_idx = torch.sort(probs, descending=True) cum = sorted_probs.cumsum(dim=-1) cutoff = cum > top_p cutoff[..., 1:] = cutoff[..., :-1].clone() cutoff[..., 0] = False # NEVER mask the top token sorted_probs = sorted_probs.masked_fill(cutoff, 0.0) probs = torch.zeros_like(probs).scatter_(1, sorted_idx, sorted_probs) probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-8) # Safe division return torch.multinomial(probs, num_samples=1) def generate( self, input_ids: torch.Tensor, price: torch.Tensor, cov: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, max_length: int = 64, **generate_kwargs ): # CRITICAL FIX: Ensure pipeline contract matches model capability num_seq = generate_kwargs.get("num_return_sequences", 1) assert num_seq == 1, "Current custom generate() only supports num_return_sequences=1" # 1. Run Encoder enc_outputs = self.t5_model.encoder( input_ids=input_ids, attention_mask=attention_mask, return_dict=True ) # 2. 🟢 Apply Adapter to Encoder Outputs enc_outputs.last_hidden_state = self.adapter( covariates=cov, encoder_hidden=enc_outputs.last_hidden_state ) # 3. Standard Decoder Loop device = input_ids.device B = input_ids.shape[0] dec_input = torch.full( (B, 1), self.t5_model.config.bos_token_id, dtype=torch.long, device=device ) outputs = [] for _ in range(max_length): out = self.t5_model( encoder_outputs=enc_outputs, decoder_input_ids=dec_input, use_cache=True, return_dict=True ) logits = out.logits[:, -1, :] logits = self._restrict_logits_to_numeric(logits) next_tok = self._sample_numeric( logits=logits, temperature=generate_kwargs.get("temperature", 1.0), top_p=generate_kwargs.get("top_p", 0.9), ) # Remove early break to ensure exactly max_length outputs # if (next_tok == self.t5_model.config.eos_token_id).all(): # break dec_input = torch.cat([dec_input, next_tok], dim=1) outputs.append(next_tok) return torch.cat(outputs, dim=1) # ============================================================================= # SECTION E: DATASET # ============================================================================= class TimeSeriesDataset(Dataset): """ LCS Upgrade: Stores raw prices, scaled input_ids/labels, and local scale. p_last is removed; reconstruction is now: decoded_token * scale. """ def __init__( self, prices: np.ndarray, covariates: np.ndarray, input_ids: np.ndarray, # LCS: scaled-price tokens labels: np.ndarray, # LCS: scaled-price tokens future_prices: np.ndarray, scale: np.ndarray, # LCS: local scale per window (replaces p_last) attention_mask: Optional[np.ndarray] = None, pad_token_id: int = 0, ): self.prices = torch.FloatTensor(prices) self.covariates = torch.FloatTensor(covariates) self.input_ids = torch.LongTensor(input_ids) self.labels = torch.LongTensor(labels) self.future_prices = torch.FloatTensor(future_prices) self.scale = torch.FloatTensor(scale) # LCS if attention_mask is None: self.attention_mask = (self.input_ids != pad_token_id).long() else: self.attention_mask = torch.LongTensor(attention_mask) def __len__(self): return self.prices.shape[0] def __getitem__(self, idx): return { "price": self.prices[idx], "cov": self.covariates[idx], "input_ids": self.input_ids[idx], "labels": self.labels[idx], "future_prices": self.future_prices[idx], "scale": self.scale[idx], # LCS (replaces p_last) "attention_mask": self.attention_mask[idx], } # ============================================================================= # SECTION F: TRAINING CONFIG & SCHEDULER # ============================================================================= @dataclass class TrainingConfig: lr: float = 1e-4 lr_end: float = 3e-5 warmup_ratio: float = 0.05 weight_decay: float = 0.01 epochs: int = 12 batch_size: int = 32 gradient_clip: float = 1.0 gradient_accumulation_steps: int = 1 mixed_precision: bool = False patience: int = 6 # Stop if no improvement after 6 epochs device: str = "cuda" if torch.cuda.is_available() else "cpu" lambda_price: float = 10000.0 # weight for price MSE term # --- NEW TOGGLE --- # True = Option A (Only Decoder + Adapter learn) # False = Option B (Encoder + Decoder + Adapter learn) freeze_encoder: bool = True save_dir: str = "./checkpoints" log_every_n_steps: int = 50 def save(self, path: str): with open(path, "w") as f: json.dump(asdict(self), f, indent=2) @classmethod def load(cls, path: str): with open(path, "r") as f: return cls(**json.load(f)) def get_custom_cosine_schedule( optimizer, num_warmup_steps: int, num_training_steps: int, lr_end_factor: float = 0.3, ): def lr_lambda(current_step: int): if current_step < num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) progress = float(current_step - num_warmup_steps) / float( max(1, num_training_steps - num_warmup_steps) ) cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) return max(lr_end_factor, cosine) return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) # ============================================================================= # SECTION G: TRAINER (HYBRID LOSS + METRICS) # ============================================================================= class CGCXTrainer: def __init__( self, model, config: TrainingConfig, normalizer: ZScoreNormalizer, train_loader: DataLoader, val_loader: Optional[DataLoader], tokenizer: SimpleChronosTokenizer, mase_scale: float = 1.0, cov_cols: Optional[List[str]] = None, ): self.model = model.to(config.device) self.config = config self.normalizer = normalizer self.train_loader = train_loader self.val_loader = val_loader self.device = config.device self.tokenizer = tokenizer assert cov_cols is not None, "cov_cols must be provided" self.cov_cols = cov_cols # Collect ALL parameters that have requires_grad=True # This includes Adapter + Decoder + (optionally) Encoder trainable_params = [] for name, p in model.named_parameters(): if p.requires_grad: trainable_params.append(p) self.optimizer = torch.optim.AdamW( trainable_params, lr=config.lr, weight_decay=config.weight_decay, ) steps_per_epoch = math.ceil(len(train_loader) / config.gradient_accumulation_steps) num_training_steps = steps_per_epoch * config.epochs self.scheduler = get_custom_cosine_schedule( self.optimizer, int(num_training_steps * config.warmup_ratio), num_training_steps, lr_end_factor=config.lr_end / config.lr # Fix Bug 5 here too ) self.scaler = GradScaler('cuda') if config.mixed_precision else None self.mase_scale = float(max(mase_scale, 1e-8)) # Precompute bin VALUES (scaled prices in [0.0, 5.0]) num_bins = tokenizer.num_bins vocab_size = model.t5_model.config.vocab_size offset = model.t5_model.config.vocab_offset bin_idx = torch.arange(num_bins, dtype=torch.float32) scaled = bin_idx / (num_bins - 1) bin_values = scaled * tokenizer.range + tokenizer.min_val # [0.0 … 5.0] full_bin_values = torch.zeros(vocab_size, dtype=torch.float32) full_bin_values[offset:offset + num_bins] = bin_values self.bin_values = full_bin_values.view(1, 1, -1).to(self.device) self.bin_values.requires_grad_(False) print("✓ Trainer initialized (LCS Mode: scaled prices in [0.0, 5.0])") def _expected_scaled_price_from_logits(self, logits: torch.Tensor) -> torch.Tensor: """Returns expected scaled-price value from logits (LCS mode).""" masked_logits = self.model._restrict_logits_to_numeric(logits) probs = torch.softmax(masked_logits, dim=-1) exp_scaled_price = (probs * self.bin_values).sum(dim=-1) return exp_scaled_price def _scaled_prices_from_labels(self, labels: torch.Tensor) -> torch.Tensor: """ Maps token IDs (labels) to their numeric scaled-price values. """ # self.bin_values is shape (1, 1, vocab_size) -> flatten to (V,) vocab_vals = self.bin_values.view(-1) # Use gather without clamping. If labels are invalid, this will crash (Good). return vocab_vals.gather(0, labels.view(-1)).view_as(labels) #CHANGE def train_epoch(self, epoch: int) -> Dict[str, float]: self.model.train() total_loss = 0.0 total_ret_mse = 0.0 total_count = 0 pad_id = self.tokenizer.pad_token_id bos_id = self.model.t5_model.config.bos_token_id self.optimizer.zero_grad() for step, batch in enumerate(self.train_loader): input_ids = batch["input_ids"].to(self.device) price = batch["price"].to(self.device) cov = batch["cov"].to(self.device) labels = batch["labels"].to(self.device) mask = batch["attention_mask"].to(self.device) cov_norm = self.normalizer.transform(cov) # --- Prepare Decoder Inputs (Teacher Forcing) --- # Shift labels right: [BOS, L1, L2, ..., Ln-1] B, H = labels.shape decoder_input_ids = torch.cat( [torch.full((B, 1), bos_id, device=self.device, dtype=torch.long), labels[:, :-1]], dim=1 ) # Build decoder mask (crucial if padding ever exists) decoder_attention_mask = (decoder_input_ids != pad_id).long() with autocast('cuda', enabled=self.config.mixed_precision): # Call model with explicit decoder inputs and NO labels outputs = self.model( input_ids=input_ids, price=price, cov=cov_norm, labels=None, decoder_input_ids=decoder_input_ids, attention_mask=mask, decoder_attention_mask=decoder_attention_mask ) logits = outputs.logits # (B, H, V) # --- 1. Numeric-Only Cross Entropy --- masked_logits = self.model._restrict_logits_to_numeric(logits) ce_loss = F.cross_entropy( masked_logits.reshape(-1, masked_logits.size(-1)), labels.reshape(-1), ignore_index=pad_id ) # --- 2. LCS Scaled-Price MSE Regression --- # Predict expected scaled price, compare with true scaled-price tokens pred_scaled = self._expected_scaled_price_from_logits(logits) true_scaled = self._scaled_prices_from_labels(labels) # Masked MSE (valid tokens only) label_mask = (labels != pad_id).float() mse_elem = (pred_scaled - true_scaled) ** 2 ret_mse_loss = (mse_elem * label_mask).sum() / (label_mask.sum() + 1e-8) # Hybrid Loss hybrid_loss = ce_loss + self.config.lambda_price * ret_mse_loss loss = hybrid_loss / self.config.gradient_accumulation_steps if self.scaler: self.scaler.scale(loss).backward() else: loss.backward() if (step + 1) % self.config.gradient_accumulation_steps == 0 or (step + 1) == len(self.train_loader): if self.scaler: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_([p for p in self.model.parameters() if p.requires_grad], self.config.gradient_clip) self.scaler.step(self.optimizer) self.scaler.update() else: # torch.nn.utils.clip_grad_norm_([p for p in self.model.parameters() if p.requires_grad], self.config.gradient_clip) # NaN guard: skip the step entirely if any gradient is NaN has_nan_grad = any( p.grad is not None and not torch.isfinite(p.grad).all() for p in self.model.parameters() if p.requires_grad ) if has_nan_grad: print(f" ⚠️ NaN gradient detected at step {step}. Skipping optimizer step.") self.optimizer.zero_grad() else: self.optimizer.step() self.scheduler.step() self.optimizer.zero_grad() # Logging (unscaled) total_loss += hybrid_loss.item() * self.config.gradient_accumulation_steps total_ret_mse += ret_mse_loss.item() * self.config.gradient_accumulation_steps total_count += 1 if step % self.config.log_every_n_steps == 0: print(f" [Epoch {epoch+1} Step {step}] Loss: {hybrid_loss.item() * self.config.gradient_accumulation_steps:.4f} (CE:{ce_loss.item():.3f} ScaledMSE:{ret_mse_loss.item():.5f})") return { "train_loss": total_loss / max(1, total_count), "train_scaled_rmse": math.sqrt(total_ret_mse / max(1, total_count)) } @torch.no_grad() def validate(self): self.model.eval() total_price_sse = 0.0 total_tokens = 0 pad_id = self.tokenizer.pad_token_id bos_id = self.model.t5_model.config.bos_token_id for batch in self.val_loader: input_ids = batch["input_ids"].to(self.device) cov = batch["cov"].to(self.device) future_prices = batch["future_prices"].to(self.device) scale = batch["scale"].to(self.device) # LCS: local scale (replaces p_last) mask = batch["attention_mask"].to(self.device) labels = batch["labels"].to(self.device) # Prepare Decoder Inputs B, H = labels.shape decoder_input_ids = torch.cat( [torch.full((B, 1), bos_id, device=self.device, dtype=torch.long), labels[:, :-1]], dim=1 ) decoder_attention_mask = (decoder_input_ids != pad_id).long() outputs = self.model( input_ids=input_ids, price=batch["price"].to(self.device), cov=self.normalizer.transform(cov), labels=None, decoder_input_ids=decoder_input_ids, attention_mask=mask, decoder_attention_mask=decoder_attention_mask ) # LCS: expected scaled price → multiply by local scale to get absolute price pred_scaled = self._expected_scaled_price_from_logits(outputs.logits) pred_prices = pred_scaled * scale.unsqueeze(1) # NO cumsum, NO exp valid_mask = (labels != pad_id) error_sq = (pred_prices - future_prices) ** 2 total_price_sse += (error_sq * valid_mask).sum().item() total_tokens += valid_mask.sum().item() avg_mse = total_price_sse / total_tokens if total_tokens > 0 else 0.0 return {"val_rmse": math.sqrt(avg_mse)} def train(self): print(f"🚀 Training Started (Max Epochs: {self.config.epochs}, Patience: {self.config.patience})") best_rmse = float("inf") patience_counter = 0 # <--- NEW: Track epochs without improvement for epoch in range(self.config.epochs): metrics = self.train_epoch(epoch) val_metrics = self.validate() #fix the bug for training missmatch #print(f"Epoch {epoch+1}: Train RMSE {metrics['train_rmse']:.4f} | Val RMSE {val_metrics['val_rmse']:.4f}") print(f"Epoch {epoch+1}: Train ScaledRMSE {metrics['train_scaled_rmse']:.4f} | Val PriceRMSE {val_metrics['val_rmse']:.4f}") # --- EARLY STOPPING LOGIC --- current_val_rmse = val_metrics["val_rmse"] if val_metrics["val_rmse"] < best_rmse: best_rmse = val_metrics["val_rmse"] patience_counter = 0 # Reset counter print(f" 🔥 New Best Model! Saving... (RMSE: {best_rmse:.4f})") Path(self.config.save_dir).mkdir(parents=True, exist_ok=True) torch.save(self.model.adapter.state_dict(), os.path.join(self.config.save_dir, "best_adapter.pt")) torch.save(self.model.t5_model.state_dict(), os.path.join(self.config.save_dir, "best_t5.pt")) self.tokenizer.save_config(os.path.join(self.config.save_dir, "tokenizer_config.json")) self.normalizer.save(self.config.save_dir) with open(os.path.join(self.config.save_dir, "model_config.json"), "w") as f: # json.dump({ "t5_base_model": "amazon/chronos-t5-base", "vocab_size": self.model.t5_model.config.vocab_size, "pad_token_id": self.model.t5_model.config.pad_token_id, "eos_token_id": self.model.t5_model.config.eos_token_id, "bos_token_id": self.model.t5_model.config.bos_token_id, "vocab_offset": self.model.t5_model.config.vocab_offset, "num_bins": self.model.t5_model.config.num_bins, "num_features": self.model.adapter.cov_encoder.projection.in_features, "d_model": self.model.d_model, "embeddings_resized": True, "freeze_encoder": self.config.freeze_encoder, "cov_cols": self.cov_cols, "target_representation": "local_scale", # Strategy 1 architecture params (needed for correct reload) "xlstm_num_layers": self.model.adapter.cov_encoder.num_layers, "xlstm_expansion_factor": self.model.adapter.cov_encoder.blocks[0].up_proj.out_features // self.model.d_model, }, f, indent=2) else: # Case B: No improvement patience_counter += 1 print(f" ⚠️ No improvement. Patience: {patience_counter}/{self.config.patience}") if patience_counter >= self.config.patience: print(f"\n⏹️ Early Stopping Triggered at Epoch {epoch+1}!") print(f" Best Validation RMSE was: {best_rmse:.4f}") break # <--- Stop the loop immediately return {"best_val_rmse": best_rmse} # ============================================================================= # SECTION H: INFERENCE WRAPPER # ============================================================================= class CGCXInference: def __init__( self, model, normalizer, tokenizer, device: Optional[str] = None, ): if device is None: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = device self.model = model.to(self.device) self.model.eval() self.normalizer = normalizer self.tokenizer = tokenizer # # @torch.no_grad() def predict( self, input_ids: torch.Tensor, price: torch.Tensor, covariates: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, num_samples: int = 20, max_length: int = 1, temperature: float = 1.0, top_p: float = 0.9, ) -> np.ndarray: cov_norm = self.normalizer.transform(covariates.to(self.device)) price_np = price.cpu().numpy() local_scale = np.mean(np.abs(price_np), axis=1, keepdims=True) local_scale = np.clip(local_scale, 1e-8, None) B = input_ids.shape[0] all_pred_scaled = [] # --- THE FIX: MAXIMUM INFERENCE CHUNKING --- # Process exactly 1 window at a time to prevent OOM for i in range(B): # Slice the current window (Batch size = 1) ids_single = input_ids[i:i+1].to(self.device) cov_single = cov_norm[i:i+1].to(self.device) mask_single = attention_mask[i:i+1].to(self.device) if attention_mask is not None else None window_samples = [] # Generate the 20 samples one-by-one to save memory! for _ in range(num_samples): gen_tokens = self.model.generate( input_ids=ids_single, price=None, cov=cov_single, attention_mask=mask_single, max_length=max_length, num_return_sequences=1, temperature=temperature, top_p=top_p ) # Decode the single token array gen_tokens_np = gen_tokens.cpu().numpy() pred_scaled_exp = self.tokenizer.decode(gen_tokens_np) if np.isnan(pred_scaled_exp).any(): pred_scaled_exp = np.nan_to_num(pred_scaled_exp, nan=1.0) window_samples.append(pred_scaled_exp) # We now have 20 samples for this 1 window. Stack them. # window_samples shape: (num_samples, 1, max_length) stacked_samples = np.stack(window_samples, axis=0) # Take the median across the samples dimension median_path = np.median(stacked_samples, axis=0) # Shape: (1, max_length) all_pred_scaled.append(median_path) # --- FORCE MEMORY CLEARING --- # Delete intermediate tensors and empty the CUDA cache # so the next chunk starts with a fresh 15GB. del ids_single, cov_single, mask_single, gen_tokens torch.cuda.empty_cache() # Reassemble all windows final_pred_scaled = np.concatenate(all_pred_scaled, axis=0) # Scale back to absolute prices pred_prices = final_pred_scaled * local_scale return pred_prices ############################################### # NEW GLOBAL TIME-BASED MULTI-ASSET LOADER ############################################### def load_multi_asset_time_splits(data_dir: str, T: int, H: int, max_files: int = 200, min_len: int = 120): train_prices_list = [] train_cov_list = [] train_labels_list = [] val_prices_list = [] val_cov_list = [] val_labels_list = [] train_scales_list = [] val_scales_list = [] test_scales_list = [] test_prices_list = [] test_cov_list = [] test_labels_list = [] files = [f for f in os.listdir(data_dir) if f.endswith(".csv")] files = files[:max_files] # LIMIT TO PROTECT RAM print(f"Processing {len(files)} tickers...") for fname in files: path = os.path.join(data_dir, fname) # ---------------------- SAFE CSV LOAD ---------------------- try: df = pd.read_csv(path) except Exception as e: print(f"⚠️ Skipping {fname}: cannot load ({e})") continue if "Date" not in df.columns: print(f"⚠️ Skipping {fname}: no Date column") continue if not {"Open","High","Low","Close"}.issubset(df.columns): print(f"⚠️ Skipping {fname}: missing OHLC") continue df["Date"] = pd.to_datetime(df["Date"], errors="coerce") df = df.dropna(subset=["Date"]) df = df.sort_values("Date").drop_duplicates(subset=["Date"]) df = df.set_index("Date") # Business-day alignment #df = df.resample("B").asfreq().ffill().bfill() df = df.resample("B").asfreq().ffill() df = df.dropna() df["price"] = df["Close"] if len(df) < min_len: continue # ------------- INDICATORS (heavy, per ticker only) ---------- df = add_technical_indicators(df) df = df.dropna() if len(df) < T + H: continue # ------------- GLOBAL TIME SPLITS (NO DATA LEAKAGE) --------- df_train = df.loc[:TRAIN_END] df_val = df.loc[TRAIN_END + pd.Timedelta(days=1): VAL_END] df_test = df.loc[VAL_END + pd.Timedelta(days=1):] # Helper: turn df → sliding windows (LCS: returns 4-tuple including scales) def process(split_df, P_list, C_list, L_list, S_list): if len(split_df) >= T + H: result = build_windows_from_df(split_df, T=T, H=H) p, c, y, s = result[0], result[1], result[2], result[3] # Only append if windows were actually generated if len(p) > 0: P_list.append(p); C_list.append(c) L_list.append(y); S_list.append(s) process(df_train, train_prices_list, train_cov_list, train_labels_list, train_scales_list) process(df_val, val_prices_list, val_cov_list, val_labels_list, val_scales_list) process(df_test, test_prices_list, test_cov_list, test_labels_list, test_scales_list) # ----------------- CONCAT WITH SAFETY CHECKS ---------------------- def safe_concat(arr_list, name): if not arr_list: raise ValueError(f"No data generated for {name} split. Check Date range/Data quality.") return np.concatenate(arr_list, axis=0) # We enforce training data exists, but val/test might technically be empty in some edge cases # (though usually we want them to exist). train_prices = safe_concat(train_prices_list, "Train") train_cov = safe_concat(train_cov_list, "Train") train_labels = safe_concat(train_labels_list, "Train") train_scales = safe_concat(train_scales_list, "Train") val_prices = safe_concat(val_prices_list, "Validation") val_cov = safe_concat(val_cov_list, "Validation") val_labels = safe_concat(val_labels_list, "Validation") val_scales = safe_concat(val_scales_list, "Validation") test_prices = safe_concat(test_prices_list, "Test") test_cov = safe_concat(test_cov_list, "Test") test_labels = safe_concat(test_labels_list, "Test") test_scales = safe_concat(test_scales_list, "Test") return ( train_prices, train_cov, train_labels, train_scales, val_prices, val_cov, val_labels, val_scales, test_prices, test_cov, test_labels, test_scales, ) # ============================================================================= # SECTION J: POST-TRAINING TEST PIPELINE (METRICS + PLOTTING) # ============================================================================= def calculate_metrics(y_true, y_pred): """Calculates comprehensive regression metrics.""" # Handle NaNs if any mask = np.isfinite(y_true) & np.isfinite(y_pred) y_true = y_true[mask] y_pred = y_pred[mask] if len(y_true) == 0: return {} mse = mean_squared_error(y_true, y_pred) mae = mean_absolute_error(y_true, y_pred) rmse = np.sqrt(mse) # Mean Absolute Percentage Error (avoid div by zero) mape = np.mean(np.abs((y_true - y_pred) / (np.abs(y_true) + 1e-8))) * 100 # R-squared r2 = r2_score(y_true, y_pred) return { "MSE": mse, "RMSE": rmse, "MAE": mae, "MAPE": f"{mape:.2f}%", "R2": r2 } ######here def load_checkpoint_with_config( checkpoint_dir: str, device: str = "cuda", strict_schema: bool = True, override_freeze_encoder: Optional[bool] = None # <--- Allows switching fine-tuning modes ): print(f"\n🔧 Loading checkpoint from {checkpoint_dir}...") tok_path = os.path.join(checkpoint_dir, "tokenizer_config.json") mod_path = os.path.join(checkpoint_dir, "model_config.json") adp_path = os.path.join(checkpoint_dir, "best_adapter.pt") if not all(os.path.exists(p) for p in [tok_path, mod_path, adp_path]): raise FileNotFoundError("Missing checkpoint files.") with open(tok_path, "r") as f: tconfig = json.load(f) with open(mod_path, "r") as f: mconfig = json.load(f) saved_cov_cols = mconfig.get("cov_cols", None) assert saved_cov_cols is not None, "Checkpoint missing 'cov_cols'" if strict_schema: saved_mode = mconfig.get("target_representation", "price") tok_mode = tconfig.get("mode", "price") assert saved_mode == "local_scale", f"Model mode '{saved_mode}' != 'local_scale'" assert tok_mode == "local_scale", f"Tokenizer mode '{tok_mode}' != 'local_scale'" assert saved_cov_cols == COV_COLS, "Global COV_COLS mismatch" tokenizer = SimpleChronosTokenizer.load_config(tok_path) normalizer = ZScoreNormalizer.load(checkpoint_dir) t5_model = AutoModelForSeq2SeqLM.from_pretrained(mconfig['t5_base_model']) if mconfig.get('embeddings_resized', False): t5_model.resize_token_embeddings(mconfig['vocab_size']) for k in ["vocab_size", "pad_token_id", "eos_token_id", "bos_token_id", "vocab_offset", "num_bins"]: setattr(t5_model.config, k, mconfig[k]) # --- THE DYNAMIC MODE SWITCHER --- if override_freeze_encoder is not None: # User is forcing a mode switch (e.g., from Standard to Full Fine-Tuning) freeze_setting = override_freeze_encoder print(f"⚠️ Overriding saved freeze state. Forcing freeze_encoder={freeze_setting}") else: # Fallback to whatever was saved in JSON (handles the legacy "freeze_backbone" key) freeze_setting = mconfig.get("freeze_encoder", mconfig.get("freeze_backbone", True)) num_layers = mconfig.get("xlstm_num_layers", 3) expansion_factor = mconfig.get("xlstm_expansion_factor", 2) model = ChronosDualPath( t5_model=t5_model, num_features=mconfig['num_features'], freeze_encoder=freeze_setting, num_layers=num_layers, # <--- LOADED FROM JSON expansion_factor=expansion_factor # <--- LOADED FROM JSON ) # Load the raw state dict raw_state_dict = torch.load(adp_path, map_location=torch.device(device)) # --- torch.compile Fix --- # Strip the "_orig_mod." prefix from the compiled weights so they match the standard model clean_state_dict = {} for key, value in raw_state_dict.items(): clean_key = key.replace("_orig_mod.", "") clean_state_dict[clean_key] = value model.adapter.load_state_dict(clean_state_dict) t5_path = os.path.join(checkpoint_dir, "best_t5.pt") if os.path.exists(t5_path): model.t5_model.load_state_dict(torch.load(t5_path, map_location=torch.device(device))) print("✓ T5 weights restored from checkpoint") else: print("⚠️ No best_t5.pt found — using pretrained T5 weights (old checkpoint)") model.to(device) model.eval() print(f"✓ Model loaded successfully (Mode: local_scale)") return model, tokenizer, normalizer, saved_cov_cols def run_test_pipeline( data_dir: str, model, normalizer, tokenizer, cov_cols: list, T: int = 64, H: int = 16, num_symbols: int = 5 ): print("\n" + "="*60) print("🚀 STARTING AUTOMATED TEST PIPELINE (LCS: Local Context Scaling)") print("="*60) all_files = [f for f in os.listdir(data_dir) if f.endswith(".csv")] if not all_files: return selected_files = random.sample(all_files, min(num_symbols, len(all_files))) infer = CGCXInference(model, normalizer, tokenizer) for fname in selected_files: symbol = fname.replace('.csv', '') print(f"\n🔍 Analyzing Symbol: {symbol}") path = os.path.join(data_dir, fname) try: df = pd.read_csv(path, parse_dates=["Date"]) df = df.sort_values("Date").drop_duplicates(subset=["Date"]).set_index("Date") #df = df.resample("B").asfreq().ffill().bfill() df = df.resample("B").asfreq().ffill() df = df.dropna() df["price"] = df["Close"] df = add_technical_indicators(df).dropna() df_test = df.loc[VAL_END + pd.Timedelta(days=1):] if len(df_test) < T + H: continue prices_test, cov_test, labels_test, scales_test = build_windows_from_df( df_test, T=T, H=H, columns=cov_cols ) if len(prices_test) == 0: continue if np.any(prices_test <= 0): print(f"Skipping {fname}: non-positive prices") continue # LCS: tokenize scaled prices (price / local_scale) local_scales = np.mean(np.abs(prices_test), axis=1, keepdims=True) # (N, 1) local_scales = np.clip(local_scales, 1e-8, None) scaled_input = prices_test / local_scales input_ids_test = np.stack([tokenizer.encode(row) for row in scaled_input]) input_ids_t = torch.LongTensor(input_ids_test) price_t = torch.FloatTensor(prices_test) cov_t = torch.FloatTensor(cov_test) # Pass Mask (Fix Bug 3) input_mask = (input_ids_t != tokenizer.pad_token_id).long() # LCS Diagnostics: token diversity (scaled prices) pad_id = tokenizer.pad_token_id unique_tokens = [len(np.unique(row[row != pad_id])) for row in input_ids_test] print(f" LCS Token Diversity: {np.mean(unique_tokens):.1f}/{T}") print(f" 🤖 Predicting {len(prices_test)} sequences...") preds = infer.predict( input_ids=input_ids_t, price=price_t, covariates=cov_t, attention_mask=input_mask, num_samples=20, temperature=1.0, max_length=H ) metrics = calculate_metrics(labels_test.flatten(), preds.flatten()) rmse = metrics.get("RMSE") print(f" 📊 All-Window RMSE: {rmse:.4f}" if rmse else " ⚠️ RMSE: N/A") except Exception as e: print(f"Error: {e}") import traceback traceback.print_exc() ####################complete visualizer############################ def run_evaluations(cursor, current_date, actuals_df, model_version="v1.0"): print("\n--- Running Daily Evaluations ---") import pandas as pd import numpy as np # 1. Get previous predictions that haven't been evaluated yet cursor.execute(""" SELECT date, entity_id, predicted_value FROM predictions WHERE model_version = %s AND date < %s ORDER BY date DESC LIMIT 400 """, (model_version, current_date)) # past_preds = cursor.fetchall() if not past_preds: print("No previous predictions to evaluate yet.") return # Sanitize actuals index and compute exact grading dates actuals_df.index = pd.to_datetime(actuals_df.index).tz_localize(None).date valid_dates = sorted([d for d in actuals_df.index if d < current_date]) if len(valid_dates) < 2: print("Not enough market data to evaluate.") return target_date = valid_dates[-1] # yesterday's close prev_date = valid_dates[-2] # day before that evals_to_insert = [] # Variables for the daily summary abs_errors_pct = [] abs_errors_dollar = [] direction_hits = 0 pred_returns = [] act_returns = [] for row in past_preds: pred_date, symbol, predicted_price = row # Check if we have the actual data for this symbol today if symbol in actuals_df.columns: # Get actual price and previous close from yfinance data try: # Assuming actuals_df has the latest closes actual_price = float(actuals_df.loc[target_date, symbol]) prev_close = float(actuals_df.loc[prev_date, symbol]) # Math error_dollar = predicted_price - actual_price error_pct = error_dollar / actual_price pred_ret = (predicted_price - prev_close) / prev_close act_ret = (actual_price - prev_close) / prev_close # Did we guess the direction correctly? direction_hit = (pred_ret > 0 and act_ret > 0) or (pred_ret < 0 and act_ret < 0) evals_to_insert.append(( target_date, symbol, predicted_price, actual_price, error_dollar, error_pct, pred_ret, act_ret, direction_hit, model_version )) # Store for summary math abs_errors_dollar.append(abs(error_dollar)) abs_errors_pct.append(abs(error_pct)) pred_returns.append(pred_ret) act_returns.append(act_ret) if direction_hit: direction_hits += 1 except Exception as e: continue # # --- NEW: Save raw actuals to the DB (Stage B) --- actuals_to_insert = [] for symbol in actuals_df.columns: try: today_actual_price = float(actuals_df.loc[target_date, symbol]) actuals_to_insert.append((target_date, symbol, today_actual_price)) except: continue if actuals_to_insert: cursor.executemany(""" INSERT INTO actuals (date, symbol, actual_value) VALUES (%s, %s, %s) ON CONFLICT DO NOTHING; """, actuals_to_insert) # 2. Save Evaluations to DB if evals_to_insert: cursor.executemany(""" INSERT INTO evaluations (date, symbol, predicted_value, actual_value, error_dollar, error_pct, predicted_return, actual_return, direction_hit, model_version) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING; """, evals_to_insert) # 3. Calculate and Save Daily Summary summary_df = pd.DataFrame({"pred": pred_returns, "act": act_returns}) # --- THE FIX: Strip Numpy types and handle NaNs safely --- try: spearman = float(summary_df.corr(method='spearman').iloc[0, 1]) if np.isnan(spearman): spearman = 0.0 except: spearman = 0.0 try: mae_dollar = float(np.nanmean(abs_errors_dollar)) if np.isnan(mae_dollar): mae_dollar = 0.0 except: mae_dollar = 0.0 try: mae_pct = float(np.nanmean(abs_errors_pct)) if np.isnan(mae_pct): mae_pct = 0.0 except: mae_pct = 0.0 dir_acc = float(direction_hits / len(evals_to_insert)) if evals_to_insert else 0.0 cursor.execute(""" INSERT INTO daily_summary (date, model_version, symbols_count, mae_dollar, mae_pct, directional_accuracy, spearman_corr) VALUES (%s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING; """, (target_date, model_version, len(evals_to_insert), mae_dollar, mae_pct, dir_acc, spearman)) print(f"✅ Evaluated {len(evals_to_insert)} symbols. MAE: {mae_pct:.2%}, Accuracy: {dir_acc:.2%}") # ============================================================================= # SECTION I: MAIN PIPELINE (MULTI-ASSET) # ============================================================================= def main(): torch.set_float32_matmul_precision('high') # Enables TF32, safe for compile torch.backends.cuda.matmul.allow_tf32 = True # Belt-and-suspenders print("Start LCS (Local Context Scaling)") DATA_DIR = "raw_data" T = 64 H = 1 epochs= 12 batch_size=16 lambda_price=10000 # NOTE: Set this to False if you want Option B (Full Fine-Tuning) # If False, LOWER the LR to 1e-5 to avoid destroying weights. freeze_encoder_setting = True lr = 1e-4 if freeze_encoder_setting else 1e-5 ######### print("Load multi-asset windows") (train_prices, train_cov, train_labels, train_scales, val_prices, val_cov, val_labels, val_scales, test_prices, test_cov, test_labels, test_scales) = load_multi_asset_time_splits(DATA_DIR, T=T, H=H) print("Tokenizer (LCS: scaled prices in [0.0, 5.0])") tokenizer = SimpleChronosTokenizer(num_bins=1024, min_val=0.0, max_val=5.0) # --- HELPER: Prepare LCS Dataset --- def prepare_dataset_tensors(prices, future_prices, scales, tokenizer_obj): """ LCS Upgrade: Encodes windows as scaled prices, not log-returns. scale = mean(abs(history)) per window. input_ids = tokenize(history / scale) label_ids = tokenize(future / scale) """ # 1. Filter invalid windows valid_mask = (np.min(prices, axis=1) > 0) & (np.min(future_prices, axis=1) > 0) if not np.all(valid_mask): drop_count = np.sum(~valid_mask) print(f"⚠️ Dropping {drop_count}/{len(prices)} windows with non-positive prices.") prices = prices[valid_mask] future_prices = future_prices[valid_mask] scales = scales[valid_mask] if len(prices) == 0: raise ValueError("All data dropped due to non-positive prices.") # 2. LCS: scale each window by its local scale scaled_input = prices / scales[:, None] # (N, T) scaled_future = future_prices / scales[:, None] # (N, H) # 3. Tokenize scaled prices input_ids = np.stack([tokenizer_obj.encode(row) for row in scaled_input]) label_ids = np.stack([tokenizer_obj.encode(row) for row in scaled_future]) # 4. Mask attention_mask = (input_ids != tokenizer_obj.pad_token_id).astype(np.int64) return input_ids, label_ids, attention_mask, scales, prices, future_prices, valid_mask tr_in, tr_lb, tr_mk, tr_sc, tr_p_safe, tr_l_safe, tr_valid = prepare_dataset_tensors(train_prices, train_labels, train_scales, tokenizer) va_in, va_lb, va_mk, va_sc, va_p_safe, va_l_safe, va_valid = prepare_dataset_tensors(val_prices, val_labels, val_scales, tokenizer) # Apply mask to covariates to maintain alignment train_cov = train_cov[tr_valid] val_cov = val_cov[va_valid] # FIX: Assert labels are valid numeric bins vocab_start = tokenizer.vocab_offset vocab_end = vocab_start + tokenizer.num_bins assert tr_lb.min() >= vocab_start and tr_lb.max() < vocab_end, "Train labels contain invalid tokens" assert va_lb.min() >= vocab_start and va_lb.max() < vocab_end, "Val labels contain invalid tokens" for special in [tokenizer.pad_token_id, tokenizer.bos_token_id, tokenizer.eos_token_id]: assert not np.any(tr_lb == special), f"Train labels contain special token {special}" assert not np.any(va_lb == special), f"Val labels contain special token {special}" print("Normalizer") normalizer = ZScoreNormalizer() normalizer.fit(torch.FloatTensor(train_cov)) config = TrainingConfig( lr=lr, epochs=epochs, batch_size=batch_size, save_dir="./checkpoints_multi", lambda_price=lambda_price, freeze_encoder=freeze_encoder_setting, mixed_precision= False ) # LCS Dataset: uses scale instead of p_last train_dataset = TimeSeriesDataset( prices=tr_p_safe, covariates=train_cov, input_ids=tr_in, labels=tr_lb, future_prices=tr_l_safe, scale=tr_sc, attention_mask=tr_mk, pad_token_id=tokenizer.pad_token_id, ) val_dataset = TimeSeriesDataset( prices=va_p_safe, covariates=val_cov, input_ids=va_in, labels=va_lb, future_prices=va_l_safe, scale=va_sc, attention_mask=va_mk, pad_token_id=tokenizer.pad_token_id, ) train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=config.batch_size) print("Load chronos-T5") model_name = "amazon/chronos-t5-base" t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_name) t5_model.resize_token_embeddings(tokenizer.vocab_size) # Apply config t5_model.config.vocab_size = tokenizer.vocab_size t5_model.config.pad_token_id = tokenizer.pad_token_id t5_model.config.eos_token_id = tokenizer.eos_token_id t5_model.config.bos_token_id = tokenizer.bos_token_id t5_model.config.vocab_offset = tokenizer.vocab_offset t5_model.config.num_bins = tokenizer.num_bins model = ChronosDualPath( t5_model=t5_model, num_features=train_cov.shape[2], freeze_encoder=config.freeze_encoder ) model.adapter = torch.compile(model.adapter, dynamic=False) trainer = CGCXTrainer( model=model, config=config, normalizer=normalizer, train_loader=train_loader, val_loader=val_loader, tokenizer=tokenizer, mase_scale=1.0, cov_cols=COV_COLS ) trainer.train() print("\n🔄 RELOADING Best Model...") reloaded_model, reloaded_tokenizer, reloaded_normalizer, reloaded_cov_cols = load_checkpoint_with_config( config.save_dir, device=config.device ) run_test_pipeline( data_dir=DATA_DIR, model=reloaded_model, normalizer=reloaded_normalizer, tokenizer=reloaded_tokenizer, cov_cols=reloaded_cov_cols, T=T, H=H, num_symbols=5 ) ####tester code#### def evaluate_and_plot_symbol(symbol: str, checkpoint_dir: str, data_dir: str, T=64, H=1): """ Loads the BEST model checkpoint and visualizes predictions for a specific symbol. """ print(f"\n🎨 VISUALIZATION: Analyzing {symbol}...") # 1. Load the BEST model (Crucial: reloads the saved state, not the last epoch) # This uses your existing load_checkpoint_with_config which correctly handles xLSTM try: model, tokenizer, normalizer, cov_cols = load_checkpoint_with_config( checkpoint_dir, device="cuda" if torch.cuda.is_available() else "cpu" ) except FileNotFoundError: print("❌ Checkpoint not found. Skipping visualization.") return infer = CGCXInference(model, normalizer, tokenizer) # 2. Load Data for the Symbol file_path = os.path.join(data_dir, f"{symbol}.csv") if not os.path.exists(file_path): print(f"❌ Data file for {symbol} not found.") return # Standard Preprocessing (Must match training) df = pd.read_csv(file_path, parse_dates=["Date"]).sort_values("Date").set_index("Date") df = df.resample("B").asfreq().ffill().dropna() df["price"] = df["Close"] df = add_technical_indicators(df).dropna() # Use only Test portion (Data the model has likely not seen) VAL_END_DATE = pd.Timestamp("2023-12-31") df_test = df.loc[VAL_END_DATE + pd.Timedelta(days=1):] if len(df_test) < T + H: print("❌ Not enough data for testing this symbol.") return # 3. Build Windows result = build_windows_from_df(df_test, T, H, columns=cov_cols) prices, cov, labels, scales = result[0], result[1], result[2], result[3] if len(prices) == 0: return # 4. LCS: Tokenize scaled input prices local_scales = np.mean(np.abs(prices), axis=1, keepdims=True) # (N, 1) local_scales = np.clip(local_scales, 1e-8, None) scaled_input = prices / local_scales input_ids = np.stack([tokenizer.encode(row) for row in scaled_input]) # Convert to Tensor input_ids_t = torch.LongTensor(input_ids) price_t = torch.FloatTensor(prices) cov_t = torch.FloatTensor(cov) mask_t = (input_ids_t != tokenizer.pad_token_id).long() # 5. Predict print(f" Generating predictions for {len(prices)} windows (Median of 20 samples)...") # --- ADD THIS WRAPPER TO PREVENT MEMORY LEAKS --- with torch.no_grad(): preds = infer.predict( input_ids=input_ids_t, price=price_t, covariates=cov_t, attention_mask=mask_t, num_samples=20, temperature=1.0, max_length=H ) # 6. Calculate Metrics metrics = calculate_metrics(labels.flatten(), preds.flatten()) print(f" 📊 RMSE: {metrics.get('RMSE', 'N/A'):.4f}") print(f" 📊 R2: {metrics.get('R2', 'N/A'):.4f}") # 7. Plotting (The new part you wanted) # Plot a random window or the last window true_curve = labels.flatten() pred_curve = preds.flatten() plt.figure(figsize=(14, 7)) x_axis = np.arange(len(true_curve)) # Plot the full continuous year of actual prices plt.plot(x_axis, true_curve, label='True Price (Entire Test Set)', color='black', linewidth=2) # Overlay the model's rolling 1-day predictions for the whole year plt.plot(x_axis, pred_curve, label='Model 1-Day Predictions', color='red', linestyle='--', linewidth=1.5, alpha=0.8) plt.title(f"Full Test Set Evaluation: {symbol} ({len(true_curve)} Trading Days)") plt.xlabel("Days in Test Set") plt.ylabel("Absolute Price ($)") plt.legend() plt.grid(True, alpha=0.3) # Save plot save_path = f"prediction_{symbol}_FULL_YEAR.png" plt.savefig(save_path) print(f" 📸 Full timeline plot saved to {save_path}") plt.close() ####tester end##### # ========================================== # 2. PREDICTION LOGIC # ========================================== def fetch_latest_data(symbol: str, save_dir: str = "daily_test"): if not os.path.exists(save_dir): os.makedirs(save_dir) ticker = yf.Ticker(symbol) df = ticker.history(period="150d") if df.empty: raise ValueError(f"Could not find data for symbol: {symbol}") df = df.reset_index() file_path = os.path.join(save_dir, f"{symbol}_latest.csv") df.to_csv(file_path, index=False) return file_path def predict_tomorrow(csv_path: str, infer, tokenizer, cov_cols, T: int = 64): device = "cuda" if torch.cuda.is_available() else "cpu" # Notice we NO LONGER load the model here. It's passed in via 'infer', 'tokenizer', and 'cov_cols' df = pd.read_csv(csv_path) date_col = "Date" if "Date" in df.columns else df.columns[0] df[date_col] = pd.to_datetime(df[date_col], utc=True).dt.tz_localize(None).dt.normalize() df = df.sort_values(date_col).set_index(date_col) df = df.resample("B").asfreq().ffill().dropna() if "Close" in df.columns: df["price"] = df["Close"] elif "close" in df.columns: df["price"] = df["close"] else: return None df = add_technical_indicators(df).dropna() df_window = df.iloc[-T:] prices_np = df_window["price"].values.astype("float32") cov_np = df_window[cov_cols].values.astype("float32") local_scale = np.mean(np.abs(prices_np)) scaled_prices = prices_np / local_scale input_ids = tokenizer.encode(scaled_prices) input_ids_t = torch.LongTensor(input_ids).unsqueeze(0).to(device) price_t = torch.FloatTensor(prices_np).unsqueeze(0).to(device) cov_t = torch.FloatTensor(cov_np).unsqueeze(0).to(device) mask_t = (input_ids_t != tokenizer.pad_token_id).long().to(device) with torch.no_grad(): pred_prices = infer.predict( input_ids=input_ids_t, price=price_t, covariates=cov_t, attention_mask=mask_t, num_samples=20, max_length=1, temperature=1.0 ) return float(pred_prices[0, 0]) # 1. Setup Modal Image with your specific libraries # Added 'lxml' because pandas.read_html requires it! image = modal.Image.debian_slim().pip_install( "torch", "pandas", "numpy", "yfinance", "requests", "lxml", "psycopg2-binary", "transformers", "scikit-learn", "matplotlib" ) app = modal.App("daily-ml-job") volume = modal.Volume.from_name("model-volume") @app.function( image=image, volumes={"/data": volume}, schedule=modal.Cron("0 11 * * 1-5"), # 8:00 AM UTC gpu="T4", secrets=[modal.Secret.from_name("neon-db-secret")], timeout=1800 # Gives your script up to 30 mins to run since yfinance can be slow ) def run_daily_inference(): import yfinance as yf import pandas as pd import numpy as np import torch import requests import io import psycopg2 from datetime import datetime, timezone print(f"⏰ SERVER WAKE-UP TIME (UTC): {datetime.now(timezone.utc)}") # ========================================== # PASTE YOUR CUSTOM CLASSES & FUNCTIONS HERE # (e.g., load_checkpoint_with_config, CGCXInference, add_technical_indicators) # ========================================== # --- HELPER FUNCTIONS --- HEADERS = {"User-Agent": "Mozilla/5.0"} def get_wiki_tickers(url, table_keywords): tickers = [] try: response = requests.get(url, headers=HEADERS) response.raise_for_status() tables = pd.read_html(io.StringIO(response.text)) target_table = next((t for t in tables if any(k in [c.lower() for c in t.columns] for k in table_keywords)), None) if target_table is not None: col_name = next(c for c in target_table.columns if c.lower() in table_keywords) tickers = target_table[col_name].tolist() except Exception as e: print(f" -> Error scraping {url}: {e}") return [str(t).replace('.', '-') for t in tickers if isinstance(t, str)] def get_all_nasdaq_tickers(): try: url = "ftp://ftp.nasdaqtrader.com/symboldirectory/nasdaqtraded.txt" df = pd.read_csv(url, sep='|') df = df[(df['Test Issue'] == 'N') & (df['NASDAQ Symbol'].notnull())] return [str(t).replace('.', '-') for t in df['Symbol'].dropna().tolist()] except Exception: return get_wiki_tickers("https://en.wikipedia.org/wiki/Nasdaq-100", ['ticker', 'symbol']) def fetch_market_caps(tickers): caps = {} for i, ticker in enumerate(tickers): try: mc = yf.Ticker(ticker).fast_info.get('market_cap') if mc: caps[ticker] = mc except: continue return caps def fetch_latest_data(symbol: str, save_dir: str = "/tmp/daily_test"): # We use /tmp because it is the standard temporary folder in Linux/Modal if not os.path.exists(save_dir): os.makedirs(save_dir) ticker = yf.Ticker(symbol) df = ticker.history(period="150d") if df.empty: raise ValueError(f"Could not find data for symbol: {symbol}") df = df.reset_index() # TEMPORAL FIREWALL: strip timezone from date index and cut today's partial row df["Date"] = pd.to_datetime(df["Date"]).dt.tz_localize(None).dt.date server_today = datetime.utcnow().date() df = df[df["Date"] < server_today] if df.empty: raise ValueError(f"No historical data for {symbol} after firewall cut.") file_path = os.path.join(save_dir, f"{symbol}_latest.csv") df.to_csv(file_path, index=False) return file_path # ========================================== # CLOUD-NATIVE LOGIC UPDATES # ========================================== # Connect to Neon to establish state print("Connecting to Neon Database...") conn = psycopg2.connect(os.environ["DATABASE_URL"]) cursor = conn.cursor() ### def get_target_symbols_from_db(): print("--- 1. Reading Golden Symbol List from Volume ---") # Read the 400 symbols from your uploaded file target_symbols_path = "/data/my_model_dir/symbols.txt" with open(target_symbols_path, "r") as f: pooled_400 = set(line.strip().upper() for line in f if line.strip()) print(f"Loaded {len(pooled_400)} target symbols from file.") print("--- 2. Checking Neon Database for existing predictions ---") # Ask Neon which ones we already processed today/historically cursor.execute("SELECT DISTINCT entity_id FROM predictions") downloaded_symbols = {row[0] for row in cursor.fetchall()} # Split them into existing and new existing_list = list(downloaded_symbols.intersection(pooled_400)) remaining_pool = pooled_400 - downloaded_symbols new_list = list(remaining_pool) return existing_list, new_list # ========================================== # MAIN EXECUTION BLOCK # ========================================== run_id = str(uuid.uuid4()) today = datetime.now().date() model_version = "v1.0" existing_symbols, new_symbols = get_target_symbols_from_db() print(f"\nTargeting {len(existing_symbols)} existing symbols and {len(new_symbols)} new symbols.") # --- NEW: Close the DB connection so it doesn't time out while doing math --- cursor.close() conn.close() print("Closed initial DB connection. Starting ML heavy lifting...") print("\n⚙️ Loading neural network model into memory...") print("\n⚙️ Loading neural network model into memory...") device = "cuda" if torch.cuda.is_available() else "cpu" # --- CLOUD UPDATE: Point to the Modal Volume Path --- checkpoint_dir = "/data/my_model_dir" model, tokenizer, normalizer, cov_cols = load_checkpoint_with_config(checkpoint_dir, device=device) infer = CGCXInference(model, normalizer, tokenizer, device=device) print("✅ Model loaded successfully!") all_predictions = [] # Combine existing and new into one loop for cleaner database insertion for sym in existing_symbols + new_symbols: try: path = fetch_latest_data(sym, save_dir="/tmp/daily_test") price_pred = predict_tomorrow(path, infer, tokenizer, cov_cols) if price_pred: # Format exactly to match the Neon 'predictions' SQL table all_predictions.append((run_id, today, sym, price_pred, model_version)) print(f"{sym}: ${price_pred:.2f}") except Exception as e: print(f"Skipping {sym} - Error: {e}") ## # --- CLOUD UPDATE: SAVE RESULTS TO DATABASE --- print(f"\n💾 Saving {len(all_predictions)} predictions to Neon Database...") if all_predictions: print("Waking up Neon Database for saving...") conn = psycopg2.connect(os.environ["DATABASE_URL"]) cursor = conn.cursor() insert_query = """ INSERT INTO predictions (run_id, date, entity_id, predicted_value, model_version) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (date, entity_id, model_version) DO NOTHING; """ cursor.executemany(insert_query, all_predictions) # --- THE FIX: Fetch a clean DataFrame of Actuals --- print("Fetching actual closing prices for evaluation...") all_syms = existing_symbols + new_symbols try: # Download the last 5 days of closing prices for all symbols raw_actuals = yf.download(all_syms, period="10d")["Close"] actuals_df = pd.DataFrame(raw_actuals) actuals_df.index = pd.to_datetime(actuals_df.index).tz_localize(None).date server_today = datetime.utcnow().date() actuals_df = actuals_df[actuals_df.index < server_today] actuals_df = actuals_df.dropna(how='all') # Run evaluations using the correct 'today' and 'actuals_df' run_evaluations(cursor, current_date=today, actuals_df=actuals_df) except Exception as e: print(f"⚠️ Warning: Evaluations failed, but predictions are safe. Error: {e}") # Finally, commit everything conn.commit() print("✅ Database commit successful. Pipeline complete.") cursor.close() conn.close() print("✅ Database save complete!") else: print("⚠️ No predictions generated today.") @app.local_entrypoint() def main(): run_daily_inference.remote()