rasmodev commited on
Commit
f111c66
·
1 Parent(s): 4e6210b
Files changed (1) hide show
  1. main.py +58 -0
main.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ import pickle
4
+ import pandas as pd
5
+
6
+ # Load the saved model
7
+ with open("model_and_key_components.pkl", "rb") as f:
8
+ components = pickle.load(f)
9
+ dt_model = components['model']
10
+
11
+ app = FastAPI()
12
+
13
+ class IncomePredictionRequest(BaseModel):
14
+ age: int
15
+ gender: str
16
+ education: str
17
+ worker_class: str
18
+ marital_status: str
19
+ race: str
20
+ is_hispanic: str
21
+ employment_commitment: str
22
+ employment_stat: int
23
+ wage_per_hour: int
24
+ working_week_per_year: int
25
+ industry_code: int
26
+ industry_code_main: str
27
+ occupation_code: int
28
+ occupation_code_main: str
29
+ total_employed: int
30
+ household_stat: str
31
+ household_summary: str
32
+ vet_benefit: int
33
+ tax_status: str
34
+ gains: int
35
+ losses: int
36
+ stocks_status: int
37
+ citizenship: str
38
+ mig_year: int
39
+ country_of_birth_own: str
40
+ importance_of_record: float
41
+
42
+
43
+ class IncomePredictionResponse(BaseModel):
44
+ income_prediction: str
45
+ prediction_probability: float
46
+
47
+ @app.post("/predict/")
48
+ async def predict_income(data: IncomePredictionRequest):
49
+ input_data = data.dict()
50
+ input_df = pd.DataFrame([input_data])
51
+ prediction = dt_model.predict(input_df)
52
+ prediction_proba = dt_model.predict_proba(input_df)
53
+ prediction_result = "Income over $50K" if prediction[0] == 1 else "Income under $50K"
54
+ return {"income_prediction": prediction_result, "prediction_probability": prediction_proba[0][1]}
55
+
56
+ if __name__ == "__main__":
57
+ import uvicorn
58
+ uvicorn.run(app, host="0.0.0.0", port=8000)