Update app.py

app.py

@@ -14,10 +14,10 @@ from sklearn.pipeline import Pipeline
 # ========== Config ==========
 DATA_PATH = "data/synthetic_transactions_samples_5000.csv"
 MODEL_DIR = "models"
-MODEL_PATH = os.path.join(MODEL_DIR, "logreg_pipeline.pkl")
+MODEL_PATH = os.path.join(MODEL_DIR, "logreg_model.pkl")
 
 # ========== FastAPI Init ==========
-app = FastAPI(title="TFIDF Logistic Regression Classifier")
+app = FastAPI()
 
 # ========== Input Schema ==========
 class TransactionData(BaseModel):
@@ -132,38 +132,46 @@ def create_text_input(row):
         Beneficial Owner: {row['Beneficial_Owner']}
     """
 
-# ========== Endpoints ==========
+# ========== Root ==========
+@app.get("/")
+def root():
+    return {"message": "TF-IDF Logistic Regression API is running."}
 
+# ========== API Routes ==========
 @app.post("/train")
 def train_model():
-    df = pd.read_csv(DATA_PATH).fillna("")
+    df = pd.read_csv(DATA_PATH)
+    df = df.fillna("")
     df["text_input"] = df.apply(create_text_input, axis=1)
 
     X = df["text_input"]
-    y = df[["Maker_Action", "Escalation_Level", "Risk_Category", "Risk_Drivers", "Investigation_Outcome"]]
+    y = df[["Maker_Action", "Escalation_Level", "Risk_Category", "Risk_Drivers", "Investigation_Outcome", "Red_Flag_Reason"]]
 
     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
+    vectorizer = TfidfVectorizer()
+    classifier = MultiOutputClassifier(LogisticRegression(max_iter=1000))
     pipeline = Pipeline([
-        ("vectorizer", TfidfVectorizer()),
-        ("classifier", MultiOutputClassifier(LogisticRegression(max_iter=1000)))
+        ("vectorizer", vectorizer),
+        ("classifier", classifier)
     ])
 
     pipeline.fit(X_train, y_train)
+
     os.makedirs(MODEL_DIR, exist_ok=True)
     joblib.dump(pipeline, MODEL_PATH)
 
     accuracy = pipeline.score(X_test, y_test)
-    return {"message": "Model trained successfully.", "accuracy": accuracy}
+    return {"message": "Model trained and saved.", "accuracy": accuracy}
 
 @app.post("/predict")
 def predict(request: TransactionData):
     try:
         model = joblib.load(MODEL_PATH)
-        input_df = pd.DataFrame([request.dict()]).fillna("")
-        text_input = create_text_input(input_df.iloc[0])
+        input_data = pd.DataFrame([request.dict()])
+        input_data = input_data.fillna("")
+        text_input = create_text_input(input_data.iloc[0])
         prediction = model.predict([text_input])[0]
-
         return {
             "Maker_Action": prediction[0],
             "Escalation_Level": prediction[1],
@@ -179,5 +187,6 @@ def validate_input(request: TransactionData):
     return {"message": "Input is valid."}
 
 @app.get("/test")
-def test():
-    return {"message": "API is working."}
+def test_api():
+    return {"message": "Test successful."}
+