Update app.py

app.py

@@ -1,17 +1,16 @@
 import gradio as gr
 import pandas as pd
-import torch
-import plotly.graph_objects as go
-from plotly.subplots import make_subplots
 import numpy as np
 import yfinance as yf
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 from timesfm import TimesFm, TimesFmHparams, TimesFmCheckpoint
 
-# --- 1. FIXED INITIALIZATION ---
-# The new API requires 'hparams' and 'checkpoint' objects, not flat arguments.
+# --- JAX INITIALIZATION ---
+# Using the native JAX/Flax checkpoint
 tfm = TimesFm(
     hparams=TimesFmHparams(
-        backend="cpu",
+        backend="cpu", # JAX handles CPU/GPU automatically
         per_core_batch_size=32,
         horizon_len=128,
         context_len=512,
@@ -19,150 +18,82 @@ tfm = TimesFm(
         model_dims=1280,
     ),
     checkpoint=TimesFmCheckpoint(
-        huggingface_repo_id="google/timesfm-1.0-200m"
+        huggingface_repo_id="google/timesfm-1.0-200m" # Original JAX weights
     ),
 )
 
 def get_financial_plot(df, forecast_df, ticker, is_backtest=False):
-    # Create Professional 2-Row Chart (Price & Volume)
     fig = make_subplots(rows=2, cols=1, shared_xaxes=True, 
                         vertical_spacing=0.05, row_heights=[0.75, 0.25])
 
-    # A. Historical Data (Blue)
-    # If backtesting, we only show 'training' data in blue to prove the AI didn't see the rest
+    # Plotting logic remains consistent for professional look
     display_df = df[:-30] if is_backtest else df
     
-    fig.add_trace(go.Scatter(
-        x=display_df.index, y=display_df['Close'],
-        name='Historical Price',
-        line=dict(color='#2962FF', width=2)
-    ), row=1, col=1)
+    # 1. Price
+    fig.add_trace(go.Scatter(x=display_df.index, y=display_df['Close'], name='History', line=dict(color='#2962FF')), row=1, col=1)
 
-    # B. The "Truth" (Grey Line - Only for Backtesting)
+    # 2. Backtest Reality (Truth)
     if is_backtest:
-        truth_df = df[-30:]
-        fig.add_trace(go.Scatter(
-            x=truth_df.index, y=truth_df['Close'],
-            name='Actual Market Move',
-            line=dict(color='#787b86', width=2, dash='dot')
-        ), row=1, col=1)
+        fig.add_trace(go.Scatter(x=df.index[-30:], y=df['Close'][-30:], name='Actual', line=dict(color='#787b86', dash='dot')), row=1, col=1)
 
-    # C. AI Forecast (Red)
-    # Connect the forecast line to the last historical point for a seamless look
-    last_hist_date = display_df.index[-1]
-    last_hist_val = display_df['Close'].iloc[-1]
-    
-    # Prepend the last historical point to the forecast data
-    fc_dates = [last_hist_date] + list(forecast_df['ds'])
-    fc_vals = [last_hist_val] + list(forecast_df['timesfm'])
-    
-    fig.add_trace(go.Scatter(
-        x=fc_dates, y=fc_vals,
-        name='AI Prediction',
-        line=dict(color='#F23645', width=3)
-    ), row=1, col=1)
+    # 3. Forecast
+    fc_dates = [display_df.index[-1]] + list(forecast_df['ds'])
+    fc_vals = [display_df['Close'].iloc[-1]] + list(forecast_df['timesfm'])
+    fig.add_trace(go.Scatter(x=fc_dates, y=fc_vals, name='AI Forecast', line=dict(color='#F23645', width=3)), row=1, col=1)
 
-    # D. Volume (Bars)
-    fig.add_trace(go.Bar(
-        x=df.index, y=df['Volume'],
-        name='Volume',
-        marker_color='rgba(38, 166, 154, 0.5)'
-    ), row=2, col=1)
+    # 4. Volume
+    fig.add_trace(go.Bar(x=df.index, y=df['Volume'], name='Volume', marker_color='rgba(38, 166, 154, 0.5)'), row=2, col=1)
 
-    # TradingView-Style Dark Theme
-    fig.update_layout(
-        template='plotly_dark',
-        paper_bgcolor='#131722',
-        plot_bgcolor='#131722',
-        margin=dict(l=10, r=10, t=40, b=10),
-        legend=dict(orientation="h", y=1.02, x=0),
-        hovermode="x unified"
-    )
-    fig.update_yaxes(gridcolor='#2a2e39', side='right')
-    fig.update_xaxes(gridcolor='#2a2e39')
-    
+    fig.update_layout(template='plotly_dark', paper_bgcolor='#131722', plot_bgcolor='#131722', margin=dict(l=10, r=10, t=40, b=10))
     return fig
 
 def run_analysis(ticker, horizon, mode):
     try:
-        # Fetch Data
-        df = yf.download(ticker, period="2y") # Get more data for stability
+        df = yf.download(ticker, period="2y")
         if df.empty: return None, "⚠️ Ticker Not Found"
         
-        # --- MODE LOGIC ---
-        if mode == "Backtest (Reality Check)":
-            # Hide the last 30 days from the AI
-            train_df = df[:-30]
-            horizon = 30 # Fixed horizon for backtest comparison
-        else:
-            # Use full data for real future prediction
-            train_df = df
+        train_df = df[:-30] if mode == "Backtest (Reality Check)" else df
         
-        # Format for TimesFM
+        # Prepare inputs for TimesFM
         input_df = pd.DataFrame({
             'unique_id': [ticker],
             'ds': train_df.index,
             'y': train_df['Close'].values
         })
 
-        # Run Forecast
+        # Forecast Execution (Native JAX backend)
         forecast_df, _ = tfm.forecast_on_df(
             inputs=input_df,
             freq="D", 
-            value_name="y",
-            forecast_context_len=512 # Explicitly use context
+            value_name="y"
         )
-        forecast_df = forecast_df.head(horizon)
+        forecast_df = forecast_df.head(horizon if mode != "Backtest (Reality Check)" else 30)
         
-        # Generate Signal (Only for Future Mode)
+        # Signal Generation
         if mode != "Backtest (Reality Check)":
-            start_price = train_df['Close'].iloc[-1]
-            end_price = forecast_df['timesfm'].iloc[-1]
-            pct_change = ((end_price - start_price) / start_price) * 100
-            
-            color = "#00ff88" if pct_change > 0 else "#ff4444"
-            direction = "BULLISH" if pct_change > 0 else "BEARISH"
-            signal_html = f"<h3 style='color: {color}; margin: 0;'>{direction} ({pct_change:+.2f}%)</h3>"
+            pct = ((forecast_df['timesfm'].iloc[-1] - train_df['Close'].iloc[-1]) / train_df['Close'].iloc[-1]) * 100
+            signal = f"<h3 style='color: {'#00ff88' if pct > 0 else '#ff4444'};'>{ 'BULLISH' if pct > 0 else 'BEARISH' } ({pct:+.2f}%)</h3>"
         else:
-            # Calculate Accuracy Score for Backtest
-            real_end = df['Close'].iloc[-1]
-            pred_end = forecast_df['timesfm'].iloc[-1]
-            error = abs(real_end - pred_end) / real_end * 100
-            accuracy = 100 - error
-            signal_html = f"<h3 style='color: #FFD700; margin: 0;'>AI Accuracy: {accuracy:.1f}%</h3>"
+            acc = 100 - (abs(df['Close'].iloc[-1] - forecast_df['timesfm'].iloc[-1]) / df['Close'].iloc[-1] * 100)
+            signal = f"<h3 style='color: #FFD700;'>Model Confidence: {acc:.1f}%</h3>"
 
-        fig = get_financial_plot(df, forecast_df, ticker, is_backtest=(mode == "Backtest (Reality Check)"))
-        
-        return fig, signal_html
+        return get_financial_plot(df, forecast_df, ticker, (mode == "Backtest (Reality Check)")), signal
         
     except Exception as e:
-        return None, f"<span style='color:red'>Error: {str(e)}</span>"
+        return None, f"Runtime Error: {str(e)}"
 
-# UI Layout
+# --- UI Layout ---
 with gr.Blocks(theme=gr.themes.Default(), css=".gradio-container {background-color: #000000}") as demo:
-    gr.HTML("""
-    <div style='background-color: #131722; padding: 15px; border-bottom: 2px solid #2962FF;'>
-        <h2 style='color: white; margin:0; text-align: center; font-family: sans-serif;'>
-            G-TIMES <span style='color: #2962FF;'>PRO TERMINAL</span>
-        </h2>
-    </div>
-    """)
-    
+    gr.HTML("<h2 style='color: #2962FF; text-align: center;'>G-TIMES JAX TERMINAL</h2>")
     with gr.Row():
-        with gr.Column(scale=1, min_width=300):
-            with gr.Group():
-                ticker_in = gr.Textbox(label="SYMBOL", value="BTC-USD")
-                mode_in = gr.Radio(["Future Forecast", "Backtest (Reality Check)"], 
-                                 label="ANALYSIS MODE", value="Future Forecast")
-                days_in = gr.Slider(5, 128, value=14, label="Horizon (Days)")
-                
-                btn = gr.Button("EXECUTE STRATEGY", variant="primary")
-            
-            gr.HTML("<br>")
-            result_box = gr.HTML(label="Signal")
-            
+        with gr.Column(scale=1):
+            ticker_in = gr.Textbox(label="SYMBOL", value="NVDA")
+            mode_in = gr.Radio(["Future Forecast", "Backtest (Reality Check)"], label="MODE", value="Future Forecast")
+            days_in = gr.Slider(5, 128, value=30, label="Days")
+            btn = gr.Button("RUN JAX INFERENCE", variant="primary")
+            result_box = gr.HTML()
         with gr.Column(scale=4):
-            plot_out = gr.Plot(label="Technical Chart")
+            plot_out = gr.Plot()
 
     btn.click(run_analysis, [ticker_in, days_in, mode_in], [plot_out, result_box])