Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,161 @@
+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()