feat: add multimodal inference service for transformer + finbert

Author: VisionExpo

backend/api/predict.py

@@ -1,26 +1,28 @@
-from fastapi import APIRouter, HTTPException
-from pydantic import BaseModel
-from backend.services import inference
-import logging
+import torch
+from fastapi import APIRouter
+
+from backend.schemas.predict import PredictRequest, PredictResponse
+from backend.services.inference import InferenceService
 
 router = APIRouter()
 
-class PredictionRequest(BaseModel):
-    ticker: str
-    date: str
+# singleton inference service
+inference_service = InferenceService(
+    model_path="models/fusion_model.pt",
+    input_dim=10,
+)
+
+
+@router.post("/predict", response_model=PredictResponse)
+def predict(req: PredictRequest):
+    market_x = torch.tensor(req.market_sequence).unsqueeze(0)
+    input_ids = torch.tensor(req.input_ids).unsqueeze(0)
+    attention_mask = torch.tensor(req.attention_mask).unsqueeze(0)
 
-class PredictionResponse(BaseModel):
-    predicted_price: float
+    out = inference_service.predict(
+        market_x,
+        input_ids,
+        attention_mask,
+    )
 
-@router.post("/predict", response_model=PredictionResponse)
-def predict(request: PredictionRequest):
-    try:
-        prediction = inference.make_prediction(ticker=request.ticker, date_str=request.date)
-        return {"predicted_price": prediction}
-    except FileNotFoundError as e:
-        raise HTTPException(status_code=404, detail=str(e))
-    except ValueError as e:
-        raise HTTPException(status_code=400, detail=str(e))
-    except Exception as e:
-        logging.error(f"Error during prediction: {e}")
-        raise HTTPException(status_code=500, detail="An internal error occurred during prediction.")
+    return PredictResponse(**out)

backend/main.py

@@ -1,50 +1,6 @@
 from fastapi import FastAPI
-from contextlib import asynccontextmanager
-from backend.services import inference
-from backend.routers import predict, backtest, drift, screener
-import logging
+from backend.api import predict
 
-logging.basicConfig(level=logging.INFO)
+app = FastAPI(title="ArthaQuant API")
 
-# --- Lifespan Events ---
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Load all ML resources on startup
-    inference.load_resources()
-    yield
-    # Clean up resources on shutdown
-    inference.cache.clear()
-
-app = FastAPI(
-    title="Stock Market AI API (Refactored)",
-    lifespan=lifespan
-)
-
-# --- Include Routers ---
-# This brings in all the endpoints from your other files
 app.include_router(predict.router)
-app.include_router(backtest.router)
-app.include_router(drift.router)
-app.include_router(screener.router)
-
-# --- Other General Endpoints ---
-from backend.services.inference import cache
-from fastapi import HTTPException
-
-@app.get("/tickers")
-def get_tickers():
-    data_df = cache.get("data")
-    if data_df is None: raise HTTPException(status_code=500, detail="Data not loaded.")
-    return data_df['Ticker'].unique().tolist()
-
-@app.get("/history/{ticker}")
-def get_history(ticker: str):
-    data_df = cache.get("data")
-    if data_df is None: raise HTTPException(status_code=500, detail="Data not loaded.")
-    ticker_data = data_df[data_df['Ticker'] == ticker.upper()]
-    if ticker_data.empty: raise HTTPException(status_code=404, detail=f"Ticker '{ticker}' not found.")
-    return ticker_data[['Date', 'Close']].to_dict('records')
-
-@app.get("/")
-def read_root():
-    return {"message": "Welcome to the Stock Market AI API!"}

backend/schemas/predict.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel
+from typing import List
+
+
+class PredictRequest(BaseModel):
+    symbol: str
+    market_sequence: List[List[float]]  # (T, F)
+    input_ids: List[int]
+    attention_mask: List[int]
+
+
+class PredictResponse(BaseModel):
+    p_up: float
+    expected_return: float
+    uncertainty: float

backend/services/inference.py

@@ -1,78 +1,43 @@
 import torch
-import numpy as np
-import pandas as pd
-from pathlib import Path
-import logging
-import joblib
-import sys
 
-# Add project root to path
-project_root = str(Path(__file__).resolve().parents[2])
-if project_root not in sys.path:
-    sys.path.insert(0, project_root)
-
-from pipelines.train_transformer_pipeline import TransformerModel
-
-# This dictionary will act as a global cache for loaded resources
-cache = {}
-
-def load_resources():
-    """Loads all necessary ML resources into the cache."""
-    if "model" in cache:
-        logging.info("Resources already loaded.")
-        return
-
-    logging.info("Loading resources for Transformer model...")
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    cache["device"] = device
-    
-    try:
-        cache["x_scaler"] = joblib.load(Path("data/sequences_sentiment/x_scaler.joblib"))
-        cache["y_scaler"] = joblib.load(Path("data/sequences_sentiment/y_scaler.joblib"))
-        data_path = Path("data/processed/final_fused_data.csv")
-        df = pd.read_csv(data_path, parse_dates=['Date'])
-        cache["data"] = df
-        
-        model_path = Path("models/transformer_v1.pt")
-        input_size = cache["x_scaler"].n_features_in_
-        
-        model = TransformerModel(input_size=input_size).to(device)
-        model.load_state_dict(torch.load(model_path, map_location=device))
-        model.eval()
-        cache["model"] = model
-        logging.info("Transformer model and all resources loaded successfully.")
-    except Exception as e:
-        logging.error(f"Failed to load resources on startup: {e}")
-
-def make_prediction(ticker: str, date_str: str) -> float:
-    """Makes a single stock prediction for a given ticker and date."""
-    if "model" not in cache:
-        raise ValueError("Model and resources are not loaded.")
-
-    model = cache["model"]
-    data_df = cache["data"]
-    x_scaler = cache["x_scaler"]
-    y_scaler = cache["y_scaler"]
-    device = cache["device"]
-
-    sequence_length = 60
-    end_date = pd.to_datetime(date_str)
-    
-    ticker_data = data_df[data_df['Ticker'] == ticker.upper()]
-    if ticker_data.empty:
-        raise FileNotFoundError(f"Data for ticker '{ticker}' not found.")
-
-    data_up_to_date = ticker_data[ticker_data['Date'] <= end_date]
-    if len(data_up_to_date) < sequence_length:
-        raise ValueError(f"Not enough historical data for {ticker} before {date_str}.")
-
-    sequence_to_predict = data_up_to_date.tail(sequence_length)
-    feature_cols = x_scaler.feature_names_in_
-    sequence_scaled = x_scaler.transform(sequence_to_predict[feature_cols])
-    
-    input_tensor = torch.from_numpy(sequence_scaled).float().unsqueeze(0).to(device)
-    with torch.no_grad():
-        prediction_scaled = model(input_tensor)
-    
-    prediction_unscaled = y_scaler.inverse_transform(prediction_scaled.cpu().numpy())[0][0]
-    return prediction_unscaled
+from ml.models.multimodal_model import MultimodalTradingModel
+
+
+class InferenceService:
+    """
+    Stateless inference service.
+    Loads model once, serves predictions.
+    """
+
+    def __init__(self, model_path: str, input_dim: int):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.model = MultimodalTradingModel(input_dim=input_dim)
+        self.model.load_state_dict(
+            torch.load(model_path, map_location=self.device)
+        )
+        self.model.to(self.device)
+        self.model.eval()
+
+    @torch.no_grad()
+    def predict(
+        self,
+        market_x: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> dict:
+        market_x = market_x.to(self.device)
+        input_ids = input_ids.to(self.device)
+        attention_mask = attention_mask.to(self.device)
+
+        output = self.model(
+            market_x,
+            input_ids,
+            attention_mask,
+        )
+
+        return {
+            "p_up": float(output["p_up"].mean().cpu()),
+            "expected_return": float(output["expected_return"].mean().cpu()),
+            "uncertainty": float(output["uncertainty"].mean().cpu()),
+        }