train.py

from pathlib import Path

import pandas as pd
import pickle
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report
from src.preprocessing import clean_text


ROOT = Path(__file__).resolve().parent
PROCESSED_DIR = ROOT / "data" / "processed"
MODELS_DIR = ROOT / "models"

balanced_path = PROCESSED_DIR / "data_balanced.csv"
fallback_path = PROCESSED_DIR / "data.csv"

# Load dataset
if balanced_path.exists():
    df = pd.read_csv(balanced_path)
elif fallback_path.exists():
    df = pd.read_csv(fallback_path)
else:
    raise FileNotFoundError(
        "No training dataset found. Expected one of: "
        f"{balanced_path} or {fallback_path}."
    )

# Clean text
df["message"] = df["message"].apply(clean_text)

# Features & labels
X = df["message"]
y = df["label"]

# TF-IDF with NGRAMS 🔥
vectorizer = TfidfVectorizer(
    max_features=10000,
    ngram_range=(1, 2)
)

X = vectorizer.fit_transform(X)

# Split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Model (improved)
model = LogisticRegression(max_iter=2000)
model.fit(X_train, y_train)

# Predict
y_pred = model.predict(X_test)

# Results
print("Accuracy:", accuracy_score(y_test, y_pred))
print(classification_report(y_test, y_pred))

# Save
MODELS_DIR.mkdir(parents=True, exist_ok=True)
with open(MODELS_DIR / "model.pkl", "wb") as model_file:
    pickle.dump(model, model_file)

with open(MODELS_DIR / "vectorizer.pkl", "wb") as vectorizer_file:
    pickle.dump(vectorizer, vectorizer_file)

print("✅ Improved model trained & saved!")