Spaces:
Sleeping
Sleeping
File size: 3,234 Bytes
d95bed5 5ade02b 3fd9376 d95bed5 643aa5a d95bed5 864bc6d 7604798 643aa5a 7604798 643aa5a 864bc6d 7604798 643aa5a 7604798 643aa5a 864bc6d 7604798 3fd9376 864bc6d 3fd9376 864bc6d 643aa5a 864bc6d 3fd9376 864bc6d 3fd9376 17a9bed 3fd9376 864bc6d 3fd9376 864bc6d 7604798 864bc6d 7604798 643aa5a 864bc6d 7604798 643aa5a 864bc6d 7604798 643aa5a 7604798 864bc6d 7604798 864bc6d 5ade02b 643aa5a | 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 | from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import os
import joblib
import pandas as pd
import threading
import time
import requests
from huggingface_hub import hf_hub_download
# --- 0. Initialize FastAPI app ---
app = FastAPI()
@app.get("/")
def read_root():
return {
"message": "Welcome to the Taxi Fare Prediction API. Use POST /predict to get predictions."
}
@app.get("/health")
def health():
return {"status": "ok"}
# --- 1. Define Input and Output Data Models ---
class InferenceInput(BaseModel):
vendor_id: str
dist_meters: float
wait_sec: float
geodetic_dist: float
mean_velocity: float
is_rush_hour: bool
model_name: str # "bog", "mex", or "uio"
class InferenceOutput(BaseModel):
trip_duration: float
model_used: str
message: str
# --- 2. Lazy-loading ML Model Manager ---
class MLModels:
def __init__(self):
self.models = {}
self.repo_id = "MatteoAld/my-ml-models" # ⚠️ Change to your actual repo
self.valid_models = {"bog", "mex", "uio"}
self.file_template = "{}_ridge_pipeline.pkl"
def get_model(self, model_name: str):
model_name = model_name.lower()
if model_name not in self.valid_models:
raise ValueError(
f"Invalid model name '{model_name}'. Choose from 'bog', 'mex', or 'uio'."
)
if model_name in self.models:
return self.models[model_name]
filename = self.file_template.format(model_name)
try:
# Download from Hugging Face Hub (cached automatically)
model_path = hf_hub_download(
repo_id=self.repo_id,
filename=filename,
cache_dir="/tmp/hf_models" # Use /tmp for write permissions
)
model = joblib.load(model_path)
self.models[model_name] = model
print(f"✅ Downloaded and loaded model '{model_name}' from Hugging Face Hub")
return model
except Exception as e:
raise RuntimeError(f"Error loading model '{model_name}': {e}")
def predict_one(self, features: dict, model_name: str):
model = self.get_model(model_name)
X_df = pd.DataFrame([features])
pred = model.predict(X_df)[0]
return float(pred)
# --- 3. Instantiate ML Model Manager ---
ml_models = MLModels()
# --- 4. Inference Endpoint ---
@app.post("/predict", response_model=InferenceOutput)
async def predict_inference(data: InferenceInput):
try:
features = data.dict()
model_name = features.pop("model_name")
trip_duration = ml_models.predict_one(features, model_name)
return InferenceOutput(
trip_duration=trip_duration,
model_used=model_name,
message=f"Inference successful using {model_name.upper()} model.",
)
except FileNotFoundError as fnf:
raise HTTPException(status_code=404, detail=str(fnf))
except ValueError as ve:
raise HTTPException(status_code=400, detail=str(ve))
except Exception as e:
raise HTTPException(status_code=500, detail=f"Inference failed: {str(e)}")
# --- 5. Run the app ---
|