Update main.py

main.py

@@ -1,7 +1,34 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Response, status, File, Form, UploadFile
+from pydantic import BaseModel
+import joblib
+from fastapi.middleware.cors import CORSMiddleware
+
+class Airflow(BaseModel):
+    topic: str
 
 app = FastAPI()
 
-@app.get("/")
+
+origins = [    "https://huggingface.co",    "https://*.huggingface.co",]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+model = joblib.load("gradient_boosting.joblib")
+
+
+@app.post("/")
 def read_root():
-    return {"Hello": "World!"}
+    return {"Hello": "World!"}
+
+
+@app.post("/api/volume_prediction")
+async def predict_volume(vol_moving_avg: float = Form(), adj_close_rolling_med: float = Form()):
+    features = [vol_moving_avg, adj_close_rolling_med]
+    predicted_volume = model.predict([features])[0]
+    return {"predicted_volume": predicted_volume}