Create train.py

train.py

@@ -0,0 +1,64 @@
+import os
+import pandas as pd
+import joblib
+from sklearn.model_selection import train_test_split
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.multioutput import MultiOutputClassifier
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import LabelEncoder
+from config import (
+    DATA_PATH, TEXT_COLUMN, LABEL_COLUMNS,
+    MODEL_SAVE_DIR, LABEL_ENCODERS_PATH,
+    TFIDF_MAX_FEATURES, NGRAM_RANGE,
+    USE_STOPWORDS, RANDOM_STATE, TEST_SIZE
+)
+
+#  Load and preprocess data
+print(" Loading dataset...")
+df = pd.read_csv(DATA_PATH)
+df.dropna(subset=[TEXT_COLUMN] + LABEL_COLUMNS, inplace=True)
+
+#  Encode each label
+label_encoders = {}
+for col in LABEL_COLUMNS:
+    le = LabelEncoder()
+    df[col] = le.fit_transform(df[col])
+    label_encoders[col] = le
+
+#  Features and targets
+X = df[TEXT_COLUMN]
+Y = df[LABEL_COLUMNS]
+
+#  Train-test split
+X_train, X_test, y_train, y_test = train_test_split(
+    X, Y, test_size=TEST_SIZE, random_state=RANDOM_STATE
+)
+
+#  Build pipeline
+stop_words = "english" if USE_STOPWORDS else None
+pipeline = Pipeline([
+    ('tfidf', TfidfVectorizer(max_features=TFIDF_MAX_FEATURES, ngram_range=NGRAM_RANGE, stop_words=stop_words)),
+    ('clf', MultiOutputClassifier(LogisticRegression(max_iter=1000, random_state=RANDOM_STATE)))
+])
+
+#  Train model
+print(" Training model...")
+pipeline.fit(X_train, y_train)
+
+#  Save full model pipeline
+model_path = os.path.join(MODEL_SAVE_DIR, "logreg_model.pkl")
+print(f" Saving model to {model_path}")
+joblib.dump(pipeline, model_path)
+
+# Save label encoders
+print(f" Saving label encoders to {LABEL_ENCODERS_PATH}")
+joblib.dump(label_encoders, LABEL_ENCODERS_PATH)
+
+#  Save TF-IDF vectorizer separately
+tfidf_vectorizer = pipeline.named_steps['tfidf']
+tfidf_path = os.path.join(MODEL_SAVE_DIR, "tfidf_vectorizer.pkl")
+print(f" Saving TF-IDF vectorizer to {tfidf_path}")
+joblib.dump(tfidf_vectorizer, tfidf_path)
+
+print("Training complete.")