Update app.py

app.py

@@ -7,32 +7,35 @@ import maven_text_preprocessing
 
 app = FastAPI()
 
-class SentimentRequest(BaseModel):
+class ClassificationRequest(BaseModel):
     message: str
     
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
 
-@app.post("/sentiment")
-def sentiment_analysis(payload: SentimentRequest):
+@app.post("/classify")
+def sentiment_analysis(payload: ClassificationRequest):
     
-    model = joblib.load("logistic_model.joblib")
+    model = joblib.load("naive_bayes.joblib")
     vectorizer = joblib.load("vectorizer.joblib")
 
-    dirty_text_series = pd.Series([payload.message])
-    clean_text = maven_text_preprocessing.clean_and_normalize(dirty_text_series)
+    clean_text = maven_text_preprocessing.clean_and_normalize(pd.Series([payload.message]))
 
     X = vectorizer.transform(clean_text)  # ⚠️ transform, NOT fit_transform
 
+    category_list = ["Politics", "Sport", "Technology", "Entertainment", "Business"]
+    
     predictions = model.predict(X)
     pred_prob = model.predict_proba(X)
 
-    pred_list = [predictions.tolist()[0],pred_prob.tolist()[0][0],pred_prob.tolist()[0][1]]
 
     return {
-        "gender": pred_list[0],
-        "f_perc": pred_list[1],
-        "m_perc": pred_list[2]
-    }
+        category_list[predictions.tolist()[0]]: pred_prob[predictions.tolist()[0]],
+        category_list[0]: pred_prob[0],
+        category_list[1]: pred_prob[1],
+        category_list[2]: pred_prob[2],
+        category_list[3]: pred_prob[3],
+        category_list[4]: pred_prob[4]
+}