File size: 1,993 Bytes
d64524a 7a6b574 d64524a 7a6b574 d64524a |
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 |
from fastapi import FastAPI
from typing import Literal
from fastapi.responses import RedirectResponse
import uvicorn
import pickle
# for data validation, so the data input by the user is realistic
from pydantic import BaseModel, Field
# request data
class Customer(BaseModel):
lead_source: Literal['organic_search', 'social_media', 'paid_ads', 'referral', 'events'] = Field(
...,
description="Source of the lead",
)
annual_income: float = Field(..., ge=0, le=109899)
number_of_courses_viewed: int = Field(..., ge=0, le=9)
# sample data
model_config = {
"json_schema_extra": {
"examples": [
{
# This dictionary below is the sample that will appear in the Swagger UI
"lead_source": "paid_ads",
"annual_income": 79276.0, # Note: Use a float (79276.0) for consistency
"number_of_courses_viewed": 2,
}
]
}
}
# response data
class PredictResponse(BaseModel):
convert_probability: float
converted: bool
app = FastAPI(title="Customer Conversion Predictor")
@app.get("/", include_in_schema=False)
def redirect_to_docs():
"""Redirects the root URL of the Space (/) to the FastAPI documentation (/docs)."""
return RedirectResponse(url="/docs")
# Load the pre-trained model
with open("model.bin", "rb") as f_in:
pipeline = pickle.load(f_in)
# Helper function to get prediction from the loaded model
def predict_single(customer_dict: dict) -> float:
return pipeline.predict_proba([customer_dict])[0, 1]
# Define the prediction endpoint
@app.post("/predict", response_model=PredictResponse)
def predict(customer: Customer):
prob = predict_single(customer.model_dump())
return PredictResponse(convert_probability=prob, converted=(prob >= 0.5))
# Run the app for local development
if __name__ == "__main__":
uvicorn.run("predict:app", host="0.0.0.0", port=9696)
|