update main

best_model_XGBoost.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6528807e265ed511558c91021f0959beca72f818c8c0328034556a2e091938fb
-size 1174163
+oid sha256:7af777eeb0d05a44845a8dcecaebdfaaf3101c93b318a1ba5f723e569e8c8b5f
+size 1174186

main.py

@@ -1,9 +1,11 @@
 import uvicorn
-import pandas as pd 
+import pandas as pd
 from pydantic import BaseModel
 from typing import List, Union
 from fastapi import FastAPI
 import joblib
+from enum import Enum
+from fastapi.responses import HTMLResponse
 
 description = """
 Welcome to the GetAround Car Value Prediction API. This app provides an endpoint to predict car values based on various features! Try it out 🕹️
@@ -30,42 +32,70 @@ app = FastAPI(
     version="0.1",
     contact={
         "name": "Antoine VERDON",
-        "email": "antoineverdon.pro@gmail.com",  # Replace with actual email
+        "email": "antoineverdon.pro@gmail.com",
     },
     openapi_tags=tags_metadata
 )
 
+class CarBrand(str, Enum):
+    citroen = "Citroën"
+    peugeot = "Peugeot"
+    pgo = "PGO"
+    renault = "Renault"
+    audi = "Audi"
+    bmw = "BMW"
+    other = "other"
+    mercedes = "Mercedes"
+    opel = "Opel"
+    volkswagen = "Volkswagen"
+    ferrari = "Ferrari"
+    maserati = "Maserati"
+    mitsubishi = "Mitsubishi"
+    nissan = "Nissan"
+    seat = "SEAT"
+    subaru = "Subaru"
+    toyota = "Toyota"
 
 class PredictionFeatures(BaseModel):
-    CarData: List[Union[str, int, bool]] = ["Renault", 193231, 85, "diesel", "black", "estate", False, True, False, False, False, False, True]
+    brand: CarBrand
+    mileage: int
+    engine_power: int
+    fuel: str
+    paint_color: str
+    car_type: str
+    private_parking_available: bool
+    has_gps: bool
+    has_air_conditioning: bool
+    automatic_car: bool
+    has_getaround_connect: bool
+    has_speed_regulator: bool
+    winter_tires: bool
 
-@app.get("/", tags=["Introduction Endpoints"])
+@app.get("/", response_class=HTMLResponse, tags=["Introduction Endpoints"])
 async def index():
     return (
-        "Hello world! This `/` is the most simple and default endpoint. If you want to learn more, check out documentation of the API at https://2nzi-getaroundapi.hf.space/docs."
+        "Hello world! This `/` is the most simple and default endpoint. "
+        "If you want to learn more, check out documentation of the API at "
+        "<a href='/docs'>/docs</a> or "
+        "<a href='https://2nzi-getaroundapi.hf.space/docs' target='_blank'>external docs</a>."
     )
-    # return "Hello world! This `/` is the most simple and default endpoint. If you want to learn more, check out documentation of the API at `/docs https://2nzi-getaroundapi.hf.space/docs`"
 
 @app.post("/predict", tags=["Machine Learning"])
 async def predict(predictionFeatures: PredictionFeatures):
     columns = [
-        'model_key', 'mileage', 'engine_power', 'fuel', 'paint_color',
+        'brand', 'mileage', 'engine_power', 'fuel', 'paint_color',
         'car_type', 'private_parking_available', 'has_gps',
         'has_air_conditioning', 'automatic_car', 'has_getaround_connect',
         'has_speed_regulator', 'winter_tires'
     ]
     
-    car_data_dict = {col: [val] for col, val in zip(columns, predictionFeatures.CarData)}
+    car_data_dict = {col: [getattr(predictionFeatures, col)] for col in columns}
     car_data = pd.DataFrame(car_data_dict)
 
-    # model_file = hf_hub_download(repo_id="2nzi/GetAround-CarPrediction", filename="best_model_XGBoost.pkl")
-    # with open(model_file, 'rb') as f:
-    #     model = pickle.load(f)
-
     model = joblib.load('best_model_XGBoost.pkl')
     prediction = model.predict(car_data)
     response = {"prediction": prediction.tolist()[0]}
     return response
 
-if __name__=="__main__":
+if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=4000)