Update app.py

app.py

@@ -15,7 +15,7 @@ LABEL_COLUMNS = [
     "Risk_Category", "Risk_Drivers", "Investigation_Outcome"
 ]
 TEXT_COLUMN = "Sanction_Context"
-MODEL_DIR = "/tmp"  # or "/data" if preferred
+MODEL_DIR = "/tmp"
 MODEL_PATH = os.path.join(MODEL_DIR, "logreg_model.pkl")
 TFIDF_PATH = os.path.join(MODEL_DIR, "tfidf_vectorizer.pkl")
 ENCODERS_PATH = os.path.join(MODEL_DIR, "label_encoders.pkl")
@@ -23,10 +23,9 @@ ENCODERS_PATH = os.path.join(MODEL_DIR, "label_encoders.pkl")
 # --- FastAPI App ---
 app = FastAPI()
 
-# --- Input Schema ---
+# --- Schemas ---
 class TransactionData(BaseModel):
     Sanction_Context: str
-    # Add all required metadata fields here if needed
 
 class PredictionRequest(BaseModel):
     transaction_data: TransactionData
@@ -44,19 +43,23 @@ def train_model(input: DataPathInput):
         df = pd.read_csv(input.data_path)
         df.dropna(subset=[TEXT_COLUMN] + LABEL_COLUMNS, inplace=True)
 
+        # Label Encoding
         label_encoders = {}
         for col in LABEL_COLUMNS:
             le = LabelEncoder()
             df[col] = le.fit_transform(df[col])
             label_encoders[col] = le
 
+        # TF-IDF
         tfidf = TfidfVectorizer(max_features=1000, ngram_range=(1, 2), stop_words="english")
         X_vec = tfidf.fit_transform(df[TEXT_COLUMN])
         y = df[LABEL_COLUMNS]
 
+        # Train Model
         model = MultiOutputClassifier(LogisticRegression(max_iter=1000))
         model.fit(X_vec, y)
 
+        # Save
         joblib.dump(model, MODEL_PATH)
         joblib.dump(tfidf, TFIDF_PATH)
         joblib.dump(label_encoders, ENCODERS_PATH)
@@ -100,7 +103,7 @@ def test_model(input: DataPathInput):
             }
             decoded_preds.append(decoded)
 
-        return {"predictions": decoded_preds[:5]}
+        return {"predictions": decoded_preds[:5]}  # limit output
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))