from fastapi import FastAPI, HTTPException from pydantic import BaseModel import pandas as pd import numpy as np import torch from pytorch_tabnet.tab_model import TabNetRegressor from contextlib import asynccontextmanager import os import requests from pathlib import Path # Define the data schema for Window 7 Conductor requests class EconRequest(BaseModel): virus_name: str econ_buffer: float # Value from 0.1 to 1.0 (IMF-WEO) population_exposure: float # Value from 100 to 100,000 (WorldPop) # Global container for the model aegis_brain = {} def download_model(): """Download the model from Hugging Face with authentication""" model_path = "aegis_window2_econ_v1.zip" model_url = "https://huggingface.co/gsstec/aegis_window2_econ_v1/resolve/main/aegis_window2_econ_v1.zip" # Get HF token from environment variable hf_token = os.getenv("HF_TOKEN") if not hf_token: print("⚠️ No HF_TOKEN environment variable found. Using public access only.") return False headers = {"Authorization": f"Bearer {hf_token}"} if not os.path.exists(model_path): print("📥 Downloading model from Hugging Face...") try: response = requests.get(model_url, headers=headers, stream=True) response.raise_for_status() with open(model_path, 'wb') as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) print("✅ Model downloaded successfully!") return True except Exception as e: print(f"❌ Failed to download model: {e}") return False else: print("✅ Model file already exists") return True @asynccontextmanager async def lifespan(app: FastAPI): # Download and load model on startup model_loaded = False if download_model(): try: model_path = "aegis_window2_econ_v1.zip" regressor = TabNetRegressor() regressor.load_model(model_path) aegis_brain["model"] = regressor model_loaded = True print("✅ TabNet Model Loaded Successfully!") except Exception as e: print(f"⚠️ Failed to load TabNet model: {e}") print("🔄 Falling back to simulation mode...") # Load risk data for cross-referencing try: rankings_url = "https://huggingface.co/gsstec/aegis_window2_econ_v1/resolve/main/Rankings.csv" hf_token = os.getenv("HF_TOKEN") if hf_token: headers = {"Authorization": f"Bearer {hf_token}"} response = requests.get(rankings_url, headers=headers) if response.status_code == 200: from io import StringIO aegis_brain["rankings"] = pd.read_csv(StringIO(response.text)) print("✅ Rankings data loaded from Hugging Face") else: raise Exception(f"Rankings CSV request failed: {response.status_code}") else: raise Exception("No HF token available") except Exception as e: print(f"⚠️ Could not load rankings from HF: {e}") # Create comprehensive mock rankings data aegis_brain["rankings"] = pd.DataFrame({ 'Virus Name': [ 'COVID-19', 'H1N1', 'SARS', 'MERS', 'Ebola', 'Zika', 'Influenza A', 'RSV', 'Marburg', 'Lassa', 'Nipah', 'Hendra' ], 'Original Score': [ 0.85, 0.65, 0.75, 0.55, 0.95, 0.45, 0.60, 0.40, 0.90, 0.70, 0.80, 0.75 ] }) print("✅ Comprehensive Mock Rankings Data Loaded") aegis_brain["model_loaded"] = model_loaded yield aegis_brain.clear() app = FastAPI(lifespan=lifespan, title="Aegis Econ API") @app.post("/predict") async def get_stability_score(data: EconRequest): model = aegis_brain.get("model") df_rankings = aegis_brain.get("rankings") model_loaded = aegis_brain.get("model_loaded", False) if df_rankings is None: raise HTTPException(status_code=500, detail="Rankings data not initialized.") # 1. Fetch risk score from Rankings.csv virus_row = df_rankings[df_rankings['Virus Name'] == data.virus_name] if virus_row.empty: # If virus not found, use a default moderate risk score base_score = 0.7 print(f"⚠️ Virus '{data.virus_name}' not found in database. Using default risk score.") else: base_score = virus_row.iloc[0]['Original Score'] risk_impact = base_score * 0.7 # 2. Predict Continental Stability if model_loaded and model is not None: # Use actual TabNet model input_vector = np.array([[base_score, risk_impact, data.econ_buffer, data.population_exposure]]) prediction = model.predict(input_vector) stability_score = float(prediction[0][0]) model_status = "trained_model" else: # Enhanced mathematical simulation econ_factor = data.econ_buffer * 0.35 population_factor = min(data.population_exposure / 100000, 1.0) * 0.35 virus_factor = (1 - base_score) * 0.30 # More sophisticated calculation stability_score = max(0.0, min(1.0, econ_factor + (1 - population_factor) + (1 - virus_factor) )) # Add realistic variance based on input parameters variance = 0.03 + (data.population_exposure / 1000000) * 0.02 stability_score += np.random.normal(0, variance) stability_score = max(0.0, min(1.0, stability_score)) model_status = "enhanced_simulation" # Determine alert level with more nuanced thresholds if stability_score < 0.25: alert_level = "CRITICAL" elif stability_score < 0.45: alert_level = "HIGH_RISK" elif stability_score < 0.65: alert_level = "MONITOR" else: alert_level = "STABLE" return { "virus": data.virus_name, "stability_score": round(stability_score, 4), "alert_level": alert_level, "model_status": model_status, "base_risk_score": round(base_score, 3), "economic_buffer_impact": round(data.econ_buffer * 0.35, 3), "population_risk_factor": round(min(data.population_exposure / 100000, 1.0), 3) } @app.get("/health") async def health(): return { "status": "operational", "hardware": "T4 GPU Active" if torch.cuda.is_available() else "CPU Mode", "model_loaded": aegis_brain.get("model_loaded", False), "available_viruses": aegis_brain.get("rankings", pd.DataFrame())['Virus Name'].tolist() if aegis_brain.get("rankings") is not None else [], "total_virus_database": len(aegis_brain.get("rankings", pd.DataFrame())), "hf_token_available": bool(os.getenv("HF_TOKEN")) } @app.get("/") async def root(): return { "message": "AEGIS Economic Stability Analysis API", "version": "1.2.0", "model_source": "https://huggingface.co/gsstec/aegis_window2_econ_v1", "features": [ "Economic stability prediction", "Multi-virus risk assessment", "Population exposure modeling", "Real-time alert classification" ], "endpoints": { "predict": "POST /predict - Get stability prediction", "health": "GET /health - System status", "root": "GET / - API information" } }