| import os, argparse, json, math, sys |
| sys.path.insert(0, '/app') |
| import torch |
| import torch.nn as nn |
| from torch.utils.data import DataLoader |
| from torch.optim import AdamW |
| from finjepa.model import FinJEPA, FinJEPALoss |
| from finjepa.data import FinancialTrajectoryDataset, build_dataloaders, generate_synthetic_data, load_hf_stock_data |
|
|
| def get_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--data_source', type=str, default='synthetic', choices=['synthetic', 'hf']) |
| parser.add_argument('--dataset_name', type=str, default='paperswithbacktest/Stocks-Daily-Price') |
| parser.add_argument('--symbols', type=str, default=None) |
| parser.add_argument('--n_assets', type=int, default=1) |
| parser.add_argument('--context_window', type=int, default=60) |
| parser.add_argument('--target_window', type=int, default=5) |
| parser.add_argument('--batch_size', type=int, default=128) |
| parser.add_argument('--embed_dim', type=int, default=128) |
| parser.add_argument('--encoder_depth', type=int, default=4) |
| parser.add_argument('--encoder_heads', type=int, default=4) |
| parser.add_argument('--predictor_depth', type=int, default=6) |
| parser.add_argument('--predictor_heads', type=int, default=4) |
| parser.add_argument('--patch_size', type=int, default=4) |
| parser.add_argument('--dropout', type=float, default=0.0) |
| parser.add_argument('--ema_decay', type=float, default=0.996) |
| parser.add_argument('--use_idm', action='store_true', default=True) |
| parser.add_argument('--lr', type=float, default=0.001) |
| parser.add_argument('--weight_decay', type=float, default=1e-6) |
| parser.add_argument('--epochs', type=int, default=50) |
| parser.add_argument('--grad_clip', type=float, default=1.0) |
| parser.add_argument('--rollout_steps', type=int, default=2) |
| parser.add_argument('--alpha', type=float, default=2.0) |
| parser.add_argument('--beta', type=float, default=1.0) |
| parser.add_argument('--delta', type=float, default=4.0) |
| parser.add_argument('--omega', type=float, default=0.5) |
| parser.add_argument('--gamma', type=float, default=0.5) |
| parser.add_argument('--output_dir', type=str, default='./outputs') |
| parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu') |
| parser.add_argument('--seed', type=int, default=42) |
| parser.add_argument('--push_to_hub', action='store_true') |
| parser.add_argument('--hub_model_id', type=str, default='ashesh8500/finjepa') |
| return parser.parse_args() |
|
|
| def set_seed(seed): |
| import random, numpy as np |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
|
|
| def train_epoch(model, loss_fn, dataloader, optimizer, device, grad_clip, rollout_steps): |
| model.train() |
| total_loss = 0 |
| total_pred = 0 |
| total_reg = 0 |
| total_temporal = 0 |
| total_idm = 0 |
| total_rollout = 0 |
| n = 0 |
| for batch in dataloader: |
| ctx = batch['context'].to(device) |
| tgt = batch['target'].to(device) |
| w = batch['weights'].to(device) |
| s = batch['signals'].to(device) |
| h = batch['hedge'].to(device) |
| outputs = model(ctx, tgt, w, s, h) |
| actions_gt = {'weights': w, 'signals': s} |
| rollout_outputs = [] |
| if rollout_steps > 1: |
| for k in range(1, rollout_steps): |
| ro = model(ctx, tgt, w, s, h) |
| rollout_outputs.append(ro) |
| loss_dict = loss_fn(outputs, actions_gt, rollout_outputs) |
| loss = loss_dict['loss'] |
| optimizer.zero_grad() |
| loss.backward() |
| if grad_clip > 0: |
| nn.utils.clip_grad_norm_(model.parameters(), grad_clip) |
| optimizer.step() |
| model.update_target() |
| total_loss += loss.item() |
| total_pred += loss_dict['loss_pred'] |
| total_reg += loss_dict['loss_reg'] |
| total_temporal += loss_dict['loss_temporal'] |
| total_idm += loss_dict['loss_idm'] |
| total_rollout += loss_dict['loss_rollout'] |
| n += 1 |
| return {'loss': total_loss / n, 'loss_pred': total_pred / n, 'loss_reg': total_reg / n, |
| 'loss_temporal': total_temporal / n, 'loss_idm': total_idm / n, 'loss_rollout': total_rollout / n} |
|
|
| @torch.no_grad() |
| def evaluate(model, loss_fn, dataloader, device, rollout_steps): |
| model.eval() |
| total_loss = 0 |
| total_pred = 0 |
| n = 0 |
| for batch in dataloader: |
| ctx = batch['context'].to(device) |
| tgt = batch['target'].to(device) |
| w = batch['weights'].to(device) |
| s = batch['signals'].to(device) |
| h = batch['hedge'].to(device) |
| outputs = model(ctx, tgt, w, s, h) |
| actions_gt = {'weights': w, 'signals': s} |
| rollout_outputs = [] |
| if rollout_steps > 1: |
| for k in range(1, rollout_steps): |
| ro = model(ctx, tgt, w, s, h) |
| rollout_outputs.append(ro) |
| loss_dict = loss_fn(outputs, actions_gt, rollout_outputs) |
| total_loss += loss_dict['loss'].item() |
| total_pred += loss_dict['loss_pred'] |
| n += 1 |
| return {'loss': total_loss / n, 'loss_pred': total_pred / n} |
|
|
| def main(): |
| args = get_args() |
| set_seed(args.seed) |
| os.makedirs(args.output_dir, exist_ok=True) |
| with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: |
| json.dump(vars(args), f, indent=2) |
| if args.data_source == 'synthetic': |
| df = generate_synthetic_data(n_timesteps=10000, n_assets=args.n_assets) |
| else: |
| symbols = args.symbols.split(',') if args.symbols else None |
| df = load_hf_stock_data(args.dataset_name, symbols=symbols) |
| print(f"Data shape: {df.shape}") |
| loaders = build_dataloaders(df, n_assets=args.n_assets, context_window=args.context_window, |
| target_window=args.target_window, batch_size=args.batch_size) |
| n_features = 14 |
| model = FinJEPA(in_features=n_features, n_assets=args.n_assets, patch_size=args.patch_size, |
| embed_dim=args.embed_dim, encoder_depth=args.encoder_depth, encoder_heads=args.encoder_heads, |
| predictor_depth=args.predictor_depth, predictor_heads=args.predictor_heads, |
| dropout=args.dropout, ema_decay=args.ema_decay, use_idm=args.use_idm).to(args.device) |
| loss_fn = FinJEPALoss(pred_loss='l1', alpha=args.alpha, beta=args.beta, delta=args.delta, |
| omega=args.omega, gamma=args.gamma).to(args.device) |
| optimizer = AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.995), weight_decay=args.weight_decay) |
| print(f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}") |
| best_val = float('inf') |
| for epoch in range(args.epochs): |
| train_metrics = train_epoch(model, loss_fn, loaders['train'], optimizer, args.device, args.grad_clip, args.rollout_steps) |
| val_metrics = evaluate(model, loss_fn, loaders['val'], args.device, args.rollout_steps) |
| print(f"Epoch {epoch+1}/{args.epochs} | train_loss={train_metrics['loss']:.4f} " |
| f"(pred={train_metrics['loss_pred']:.4f} reg={train_metrics['loss_reg']:.4f} " |
| f"temp={train_metrics['loss_temporal']:.4f} idm={train_metrics['loss_idm']:.4f} " |
| f"rollout={train_metrics['loss_rollout']:.4f}) | val_loss={val_metrics['loss']:.4f}") |
| if val_metrics['loss'] < best_val: |
| best_val = val_metrics['loss'] |
| torch.save(model.state_dict(), os.path.join(args.output_dir, 'best_model.pt')) |
| torch.save(model.state_dict(), os.path.join(args.output_dir, 'final_model.pt')) |
| if args.push_to_hub: |
| from huggingface_hub import HfApi |
| api = HfApi() |
| api.create_repo(args.hub_model_id, repo_type="model", exist_ok=True) |
| api.upload_folder(folder_path=args.output_dir, repo_id=args.hub_model_id, repo_type="model") |
| print(f"Pushed to https://huggingface.co/{args.hub_model_id}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
