Update UI with Technical Report and Metrics

app.py

@@ -1,3 +1,7 @@
+"""
+Gradio app for HFT Sequence Modeling.
+Includes overview, ablation summary, code display.
+"""
 import gradio as gr
 import torch
 import numpy as np
@@ -7,37 +11,48 @@ import os
 import sys
 from datasets import load_dataset
 from src.models import get_model
-from src.engine import quantize_model
+from src.engine import quantize_model, get_model_size
 
-# --- CONFIGURATION ---
-DATASET_NAME = "aayushkrm/wunder-fund-hft-data" 
+# ============== CONFIGURATION ==============
+DATASET_NAME = "aayushkrm/wunder-fund-hft-data"
+SAMPLE_DATA_LENGTH = 200 # Sample sequence length
+LOAD_SAMPLE = True # Fast prototype by loading only first 1% of the data, else LOAD = False
+BEST_MODEL_STRATEGY = "Strategy ED"
+BEST_MODEL_DESC = "1. SE-Mish-DeepResGRU"
+# ============== DATA LOADING ==============
+print("Initializing...")
 
-# --- LOAD DATASET ---
-print("Initializing App...")
 try:
-    print("Loading FULL dataset from Hugging Face...")
-    # FIX: Load "train" split entirely (141MB fits easily in memory)
-    # This ensures we get ALL sequence IDs, not just the first one.
-    dataset = load_dataset(DATASET_NAME, split="train")
-    df = dataset.to_pandas()
-    
-    # --- TYPE CONVERSION ---
-    df['seq_ix'] = df['seq_ix'].astype(int)
-    
-    # Get all unique IDs and sort them
-    unique_ids_np = df['seq_ix'].unique()
-    SEQ_IDS = sorted([int(x) for x in unique_ids_np.tolist()])
-    
-    print(f"✅ Loaded {len(df)} rows.")
-    print(f"✅ Found {len(SEQ_IDS)} unique sequences: {SEQ_IDS[:5]}...")
-    
+    print("Loading dataset from Hugging Face...")
+    # Load first 1% just for demo speed (Adjusted for the same error)
+    if LOAD_SAMPLE:
+        print("Loading sample dataset")
+        df = load_dataset(DATASET_NAME, split="train[:1%]").to_pandas()
+        df['seq_ix'] = df['seq_ix'].astype(int)
+        SEQ_IDS = sorted([int(x) for x in df['seq_ix'].unique().tolist()])
+    else:
+        print("Loading Full dataset")
+        df = load_dataset(DATASET_NAME, split="train").to_pandas()
+        df['seq_ix'] = df['seq_ix'].astype(int)
+        SEQ_IDS = sorted([int(x) for x in df['seq_ix'].unique().tolist()])
+        
 except Exception as e:
     print(f"⚠️ Could not load HF dataset: {e}")
     df = None
-    SEQ_IDS = [0]
+    SEQ_IDS = [0, 1, 2] # DUMMY
+    
+# Load default dataframe for visualization
+df_display = pd.DataFrame({
+    'Time Step': list(range(100)),
+    'Actual Feature 0': np.random.randn(100),
+    'Predicted Feature 0': np.random.randn(100)
+})
+
+print("Building UI...")
 
-# --- CACHED MODEL LOADER ---
+# ============== LOAD MODEL ==============
 def load_cached_model():
+    # Load saved scaler
     model = get_model("winner", input_size=32, hidden_size=256, layers=6)
     model_path = "artifacts/best_model.pt"
     
@@ -45,49 +60,60 @@ def load_cached_model():
         try:
             print(f"Loading weights from {model_path}...")
             state = torch.load(model_path, map_location='cpu')
+            # Verify dictionary load and convert to float32
             state = {k: v.float() for k, v in state.items()}
             model.load_state_dict(state)
-            print("✅ Loaded best_model.pt")
+            print(f"✅ Loaded Rank 28 solution with structure. Key file:")
+            # Force correct behavior to load back quantized models
         except Exception as e:
             print(f"⚠️ Error loading model: {e}")
     else:
         print("⚠️ Model file not found, using random weights.")
-    
     model = quantize_model(model)
     return model
-
 MODEL = load_cached_model()
 
-def inference(seq_id_input, steps_input):
-    seq_id = int(seq_id_input)
-    steps_to_plot = int(steps_input)
-    
+# --- Data for Overview/Ablation study section ----
+ablation_data = """
+| Rank | Approach | Score |
+|---|---|---|
+| 28 | Strategy ED SE-Mish-Swarm | 0.3873 |
+| N/A | Strategy DS Mish-GRU Massive | 0.3871 |
+"""
+
+# --- Load data to create and download a table of models and score ---
+def to_csv(data):
+    return data.to_csv(index=False)
+
+df_log = pd.DataFrame([
+            {"Architecture":"SE-Mish-GRU", "Technique":"INT8","Size (MB)":9.9, "Score": 0.3873 },
+            {"Architecture":"WaveNet", "Technique":"TimeOut", "Size (MB)": 23.0 , "Score": 0.0},
+        ])
+ 
+csv_text = to_csv(df_log)
+
+def plot_forecast(seq_id, steps_to_plot):
     if df is not None:
         seq_data = df[df['seq_ix'] == seq_id].sort_values('step_in_seq')
-        if len(seq_data) == 0:
-            raw_values = np.random.randn(1000, 32).astype(np.float32)
-        else:
-            raw_values = seq_data[[str(i) for i in range(32)]].values.astype(np.float32)
-            
+        raw_values = seq_data[[str(i) for i in range(32)]].values.astype(np.float32)
         mean = raw_values.mean(axis=0)
         std = raw_values.std(axis=0) + 1e-6
         norm_values = (raw_values - mean) / std
     else:
         norm_values = np.random.randn(1000, 32).astype(np.float32)
-
+            
     x = torch.tensor(norm_values).unsqueeze(0)
-    
     with torch.no_grad():
         preds = []
         h = None
         for t in range(min(len(x[0]), steps_to_plot)):
             xt = x[:, t:t+1, :]
             o, h = MODEL(xt, h)
-            preds.append(float(o.numpy()[0,0,0]))
-            
+            preds.append(o.numpy()[0,0,0])
+                
     fig = go.Figure()
     y_actual = [float(v) for v in norm_values[:steps_to_plot, 0].flatten()]
-    y_pred = preds
+    y_pred = [float(v) for v in preds]
     x_axis = list(range(len(y_actual)))
     
     fig.add_trace(go.Scatter(x=x_axis, y=y_actual, mode='lines', name='Actual', line=dict(color='gray')))
@@ -102,24 +128,129 @@ def inference(seq_id_input, steps_input):
     )
     return fig
 
-# --- UI ---
-with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
-    gr.Markdown("# ⚡ Quant-Lab: HFT Sequence Modeling")
-    gr.Markdown(f"**Strategy:** SE-Mish-DeepResGRU (Rank 28 Implementation)")
+def get_code():
+    return inspect.getsource(PredictionModel.predict)
+
+def json_code(spec):
+    if spec is None: return 0
+    try: return json.dumps(eval(spec))
+    except: return json_code(None)
+
+def main():
+    # --- UI ---
+    # Add a Theme to the app
+    with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
+        with gr.Tab("Overview"):
+            gr.Markdown(
+                """
+                ## High-Frequency Trading Sequence Modeling
+                This project tackles the problem of predicting future market states from sequences of past states. 
+                *   **Dataset:** Anonymized features resembling real-world market data.
+                *   **Objective:** Predict the next market state (32 features) given a history of 1000 steps.
+                *   **Constraint:** Solutions must be deployable within a 20MB storage limit and run efficiently on CPU hardware.
+                """
+            )
+            gr.Markdown("### Key Results (Private Leaderboard)")
+            gr.Markdown(ablation_data)
+            
+            # Data for Downloading Artifacts log as CSV
+            gr.File(value=csv_text, label="Download the Models Benchmark as CSV", file_name="models_report.csv")
     
-    with gr.Row():
-        seq_selector = gr.Dropdown(
-            choices=SEQ_IDS, 
-            label="Select Sequence ID", 
-            value=SEQ_IDS[0]
-        )
-        step_slider = gr.Slider(
-            minimum=50, maximum=1000, value=200, label="Steps"
-        )
+        with gr.Tab("Inference"):
+            gr.Markdown(f"### Model: {BEST_MODEL_STRATEGY}")
+            gr.Markdown(f"Description: {BEST_MODEL_DESC}")
+            
+            with gr.Row():
+                seq_selector = gr.Dropdown(
+                    choices=[int(x) for x in SEQ_IDS], 
+                    label="Select Market Sequence", 
+                    value=int(SEQ_IDS[0])
+                )
+                step_slider = gr.Slider(
+                    minimum=50, maximum=1000, value=SAMPLE_DATA_LENGTH, label="Steps"
+                )
     
-    plot = gr.Plot(label="Forecast")
-    btn = gr.Button("Run Inference", variant="primary")
-    btn.click(inference, inputs=[seq_selector, step_slider], outputs=plot)
+            plot = gr.Plot(label="Forecast")
+            btn = gr.Button("Run Inference", variant="primary")
+            btn.click(plot_forecast, inputs=[seq_selector, step_slider], outputs=plot)
+
+        with gr.Tab("Code"):
+            gr.Code(value="""
+import numpy as np
+from pathlib import Path
+import torch
+import json
+class PredictionModel:
+    def __init__(self):
+        torch.set_num_threads(1)
+        base = Path(__file__).parent
+        with open(base / "artifacts/config.json", "r") as f:
+            config = json.load(f)
+        
+        scaler = np.load(base / "artifacts/scaler.npz")
+        self.mean = scaler["mean"].astype(np.float32)
+        self.std = scaler["std"].astype(np.float32)
+        self.std[self.std < 1e-6] = 1e-6
+        
+        # Same R²-based weighting as Solution 6
+        self.weights = np.array(config['weights'], dtype=np.float32)
+        
+        self.device = torch.device("cpu")
+        self.models = []
+        
+        for i in range(config['num_models']):
+            model = Solution6GRU(config['input_size'], config['hidden_size'],
+                                config['num_layers'], config['dropout'])
+            state_dict = torch.load(artifacts_dir / f"model_{i}.pt", map_location=self.device)
+            model.load_state_dict(state_dict)
+            model.to(self.device)
+            model.eval()
+            self.models.append(model)
+        
+        self.current_seq_ix: Optional[int] = None
+        self.hidden_states: list = []
+    
+    def _reset_state(self) -> None:
+        self.hidden_states = [None] * len(self.models)
+    
+    def predict(self, data_point) -> np.ndarray | None:
+        if self.current_seq_ix != data_point.seq_ix:
+            self.current_seq_ix = data_point.seq_ix
+            self._reset_state()
+        
+        state_arr = np.asarray(data_point.state, dtype=np.float32)
+        scaled = (state_arr - self.mean) / self.std
+        input_tensor = torch.from_numpy(scaled).view(1, 1, -1).to(self.device)
+        
+        preds = []
+        with torch.no_grad():
+            for i, model in enumerate(self.models):
+                output, new_hidden = model(input_tensor, self.hidden_states[i])
+                self.hidden_states[i] = new_hidden.detach()
+                preds.append(output[0, 0, :].cpu().numpy())
+        
+        if not data_point.need_prediction:
+            return None
+        
+        # WEIGHTED average using R²-based weights (Solution 6 method)
+        preds_array = np.array(preds)
+        prediction = np.sum(preds_array * self.weights[:, np.newaxis], axis=0)
+        prediction = prediction * self.std + self.mean
+        return prediction.astype(np.float32)
+""", language="python")
+
+    
+    with gr.Accordion("Notes"):
+        gr.Markdown(
+            """
+            ## Further Improvements
+            1.  More sophisticated ensembling techniques. 
+            2.  Dynamic code generation for hardware optimization.
+            """
+        )
+            
+    with gr.Footer():
+            gr.Markdown("Created by Aayush Kumar. [GitHub Repository](https://github.com/aayushkrm/efficient-neural-hft)")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=False)