File size: 4,826 Bytes
3bce488 ce3d808 3bce488 ce3d808 3bce488 ce3d808 3bce488 ce3d808 3bce488 ce3d808 3bce488 ce3d808 3bce488 ce3d808 3bce488 ce3d808 3bce488 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
from fastapi import FastAPI, HTTPException, UploadFile, File
from pydantic import BaseModel
import joblib
import pandas as pd
import numpy as np
import os
from typing import List
app = FastAPI(title="Stock Prediction API", version="1.0.0")
# Global variables to store models
models = {}
class PredictionInput(BaseModel):
symbol: str = "AAPL"
sma_20: float
sma_50: float
rsi: float
macd: float
class PredictionOutput(BaseModel):
prediction: float
model_type: str
symbol: str
@app.on_event("startup")
def load_models():
"""Load models on startup."""
from pathlib import Path
BASE_DIR = Path(__file__).resolve().parent.parent.parent
model_dir = BASE_DIR / "models"
print(f"Loading models from: {model_dir}")
if not model_dir.exists():
print(f"Models directory not found at {model_dir}")
return
# iterate over subdirs (symbols)
for symbol_dir in model_dir.iterdir():
if symbol_dir.is_dir():
symbol = symbol_dir.name
print(f"Found symbol directory: {symbol}")
# Load Regression
reg_path = symbol_dir / "regression_model.pkl"
if reg_path.exists():
try:
key = f"regression_{symbol}"
models[key] = joblib.load(reg_path)
print(f"Loaded {key} from {reg_path}")
# Keep legacy 'regression' key pointing to AAPL for backward compat
if 'regression' not in models or symbol == "AAPL":
models['regression'] = models[key]
except Exception as e:
print(f"Failed to load {reg_path}: {e}")
# Load Classification
clf_path = symbol_dir / "classification_model.pkl"
if clf_path.exists():
try:
key = f"classification_{symbol}"
models[key] = joblib.load(clf_path)
print(f"Loaded {key} from {clf_path}")
if 'classification' not in models or symbol == "AAPL":
models['classification'] = models[key]
except Exception as e:
print(f"Failed to load {clf_path}: {e}")
@app.get("/health")
def health_check():
return {"status": "healthy", "models_loaded": list(models.keys())}
@app.post("/predict/price", response_model=PredictionOutput)
def predict_price(input_data: PredictionInput):
symbol = input_data.symbol
model_key = f"regression_{symbol}"
# Fallback to generic 'regression' if specific symbol not found
if model_key not in models:
if 'regression' in models:
model_key = 'regression'
else:
raise HTTPException(status_code=503, detail=f"Regression model for {symbol} not loaded")
features = [[input_data.sma_20, input_data.sma_50, input_data.rsi, input_data.macd]]
prediction = models[model_key].predict(features)[0]
return {"prediction": prediction, "model_type": str(type(models[model_key])), "symbol": symbol}
@app.post("/predict/direction", response_model=PredictionOutput)
def predict_direction(input_data: PredictionInput):
symbol = input_data.symbol
model_key = f"classification_{symbol}"
if model_key not in models:
if 'classification' in models:
model_key = 'classification'
else:
raise HTTPException(status_code=503, detail=f"Classification model for {symbol} not loaded")
features = [[input_data.sma_20, input_data.sma_50, input_data.rsi, input_data.macd]]
prediction = models[model_key].predict(features)[0]
return {"prediction": float(prediction), "model_type": str(type(models[model_key])), "symbol": symbol}
@app.post("/predict/batch")
async def predict_batch(file: UploadFile = File(...)):
if 'regression' not in models:
raise HTTPException(status_code=503, detail="Regression model not loaded")
try:
df = pd.read_csv(file.file)
required_cols = ['sma_20', 'sma_50', 'rsi', 'macd']
if not all(col in df.columns for col in required_cols):
raise HTTPException(status_code=400, detail=f"CSV must contain columns: {required_cols}")
features = df[required_cols]
predictions = models['regression'].predict(features)
results = df.copy()
results['predicted_price'] = predictions
# Return as JSON records
return results.to_dict(orient="records")
except Exception as e:
raise HTTPException(status_code=500, detail=f"Batch processing failed: {e}")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
|