Upload run_training_fast.py with huggingface_hub

run_training_fast.py

@@ -0,0 +1,59 @@
+import sys
+sys.path.insert(0, '/app')
+import torch
+from finjepa.model import FinJEPA, FinJEPALoss
+from finjepa.data import generate_synthetic_data, build_dataloaders
+from torch.optim import AdamW
+from finjepa.planner import CEMPlanner
+
+DEVICE = 'cpu'
+df = generate_synthetic_data(n_timesteps=5000, n_assets=1)
+loaders = build_dataloaders(df, n_assets=1, context_window=60, target_window=10, batch_size=64)
+
+model = FinJEPA(
+    in_features=14, n_assets=1, patch_size=4, embed_dim=64,
+    encoder_depth=2, encoder_heads=4, predictor_depth=3, predictor_heads=4,
+    ema_decay=0.996, use_idm=True,
+).to(DEVICE)
+
+loss_fn = FinJEPALoss(pred_loss='l1', alpha=2.0, beta=1.0, delta=4.0, omega=0.5, gamma=0.5).to(DEVICE)
+optimizer = AdamW(model.parameters(), lr=0.001, betas=(0.9, 0.995), weight_decay=1e-5)
+
+print(f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
+
+for epoch in range(3):
+    model.train()
+    total_loss = 0
+    n = 0
+    for batch in loaders['train']:
+        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)
+        out = model(ctx, tgt, w, s, h)
+        actions_gt = {'weights': w, 'signals': s}
+        loss_dict = loss_fn(out, actions_gt)
+        loss = loss_dict['loss']
+        optimizer.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        model.update_target()
+        total_loss += loss.item()
+        n += 1
+    print(f"Epoch {epoch+1}/3 | loss={total_loss / n:.3f}")
+
+print("Training complete. Testing planner...")
+planner = CEMPlanner(model, n_assets=1, horizon=5, n_candidates=20, n_elites=5, n_iterations=2)
+batch = next(iter(loaders['test']))
+ctx = batch['context'][0:1].to(DEVICE)
+result = planner.plan(ctx)
+print(f"Best weights: {result['weights'].detach().cpu().numpy().round(3)}")
+print(f"Best signals: {result['signals'].detach().cpu().numpy()}")
+print(f"Expected cost: {result['expected_cost'].item():.4f}")
+
+import os
+os.makedirs('/app/finjepa/outputs', exist_ok=True)
+torch.save(model.state_dict(), '/app/finjepa/outputs/fast_model.pt')
+print("Model saved.")