Merge branch 'main' of https://github.com/adhit-420/FDS

Author: adhit-420

backend/app/api/models/phish_model.py

@@ -3,8 +3,8 @@
 
 class RFModel:
     def __init__(self):
-        self.model = self.load_model("phish_model.pkl")
-        self.vectorizer = self.load_model("vectorizer.pkl")
+        self.model = self.load_model("./app/api/models/rf_model.pkl")
+        self.vectorizer = self.load_model("./app/api/models/vectorizer.pkl")
 
     def load_model(self, model_path):
         with open(model_path, "rb") as f:
@@ -14,5 +14,10 @@ def load_model(self, model_path):
     def predict(self, input_data):
         # Use the loaded model to make predictions on the input data
         input_vectorized = self.vectorizer.transform(input_data)
-        predictions = self.model.predict(input_data)
+        predictions = self.model.predict(input_vectorized)
+        return predictions
+
+    def predict_proba(self, input_data):
+        input_vectorized = self.vectorizer.transform(input_data)
+        predictions = self.model.predict_proba(input_vectorized)
         return predictions

backend/app/api/routers/model_serve.py

@@ -20,6 +20,8 @@ class OutputData(BaseModel):
 @router.get("/phish_model")
 async def predict_phish_model(url: str):
     model = RFModel()
-    pred = model.predict([url])
-    proba = model.predict_proba([url])
-    return {pred: pred, proba: proba}
+    pred = model.predict([url])[0]
+    proba = max(model.predict_proba([url])[0])
+    print(pred)
+    print(proba)
+    return {"pred": pred, "proba": proba}