Update app.py

app.py

@@ -7,7 +7,8 @@ import os
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.preprocessing import LabelEncoder
 from sklearn.multioutput import MultiOutputClassifier
-
+import config
+from sklearn.metrics import accuracy_score
 # --- Configuration ---
 LABEL_COLUMNS = [
     "Red_Flag_Reason", "Maker_Action", "Escalation_Level",
@@ -103,40 +104,46 @@ def health_check():
 @app.post("/train")
 def train():
     try:
-        os.makedirs(config.MODEL_SAVE_DIR, exist_ok=True)
+        # Load data
         df = pd.read_csv(config.DATA_PATH)
 
-        # Features and Labels
+        # Prepare features and labels
         X = df[config.TEXT_COLUMN]
         y = df[config.LABEL_COLUMNS]
 
         # Split the data
         X_train, X_test, y_train, y_test = train_test_split(
-            X, y, test_size=0.2, random_state=42
+            X, y, test_size=config.TEST_SIZE, random_state=config.RANDOM_STATE
         )
 
-        # TF-IDF vectorization
-        vectorizer = TfidfVectorizer()
+        # TF-IDF vectorizer
+        vectorizer = TfidfVectorizer(
+            max_features=config.TFIDF_MAX_FEATURES,
+            ngram_range=config.NGRAM_RANGE,
+            stop_words='english' if config.USE_STOPWORDS else None
+        )
         X_train_vec = vectorizer.fit_transform(X_train)
         X_test_vec = vectorizer.transform(X_test)
 
-        # Train MultiOutput Logistic Regression
+        # Train model
         model = MultiOutputClassifier(LogisticRegression(max_iter=1000))
         model.fit(X_train_vec, y_train)
 
-        # Predict and evaluate
+        # Predict on test data
         y_pred = model.predict(X_test_vec)
+
+        # Calculate accuracy for each label
         accuracy = {
-            col: accuracy_score(y_test[col], [pred[i] for pred in y_pred])
-            for i, col in enumerate(y.columns)
+            label: accuracy_score(y_test[label], [pred[i] for pred in y_pred])
+            for i, label in enumerate(config.LABEL_COLUMNS)
         }
 
         # Save model and vectorizer
         joblib.dump(model, config.MODEL_PATH)
-        joblib.dump(vectorizer, config.TFIDF_PATH)
+        joblib.dump(vectorizer, config.TFIDF_VECTORIZER_PATH)
 
         return {
-            "message": " Training completed successfully.",
+            "message": "Training completed successfully.",
             "accuracy": accuracy
         }