hagar-sayed/risk_tolerance

1from fastapi import FastAPI
2import joblib
3import pandas as pd
4from pydantic import BaseModel
5
6app = FastAPI()
7
8# Load the model
9model = joblib.load("risk_tolerance.joblib")
10
11# Define the feature names
12feature_names = [
13    'age', 'education_level', 'married_state', 'no_of_kids', 'life_statge',
14    'occupational_category', 'income', 'risk', 'eager'
15]
16
17# Create a Pydantic model for the input data
18class InputData(BaseModel):
19    age: int
20    education_level: int
21    married_state: int
22    no_of_kids: int
23    life_statge: int
24    occupational_category: int
25    income: float
26    risk: int
27    eager: int
28
29def categorize(value):
30    threshold_low = 33.3
31    threshold_high = 66.6
32    if value <= threshold_low:
33        return 'Low'
34    elif value <= threshold_high:
35        return 'Medium'
36    else:
37        return 'High'
38
39@app.post("/predict")
40def predict(data: InputData):
41    features = [data.age, data.education_level, data.married_state, data.no_of_kids, data.life_statge, data.occupational_category, data.income, data.risk, data.eager]
42    features_df = pd.DataFrame([features], columns=feature_names)
43    expense = model.predict(features_df)
44    return {"risk_tolerance": categorize(expense[0])}
45