File size: 1,685 Bytes
a6b5f1d | 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 | from fastapi import FastAPI
from typing import Literal
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": [
{
"lead_source": "paid_ads",
"annual_income": 79276.0,
"number_of_courses_viewed": 2,
}
]
}
}
# response data
class PredictResponse(BaseModel):
convert_probability: float
converted: bool
app = FastAPI(title="Customer Conversion Predictor")
# 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)
return PredictResponse(convert_probability=prob, converted=(prob >= 0.5))
# Run the app for local development
# The '__main__' block is for local development and will NOT run on Hugging Face Spaces
# if __name__ == "__main__":
# uvicorn.run("predict:app", host="0.0.0.0", port=9696)
|