ai-backend/model_service.py

import warnings

import pandas as pd
from sklearn.calibration import CalibratedClassifierCV
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    log_loss,
)
from sklearn.model_selection import StratifiedGroupKFold
from sklearn.pipeline import FeatureUnion, Pipeline


WORD_NGRAM_RANGE = (1, 3)
CHAR_NGRAM_RANGE = (3, 5)
MAX_FEATURES_PER_ANALYZER = 50_000
MAX_ITERATIONS = 2_000
CALIBRATION_FOLDS = 2
EVALUATION_FOLDS = 5
RANDOM_STATE = 42
MIN_RECOMMENDED_CLASS_SIZE = 100


def build_base_model():
    features = FeatureUnion(
        [
            (
                "word",
                TfidfVectorizer(
                    ngram_range=WORD_NGRAM_RANGE,
                    sublinear_tf=True,
                    min_df=2,
                    max_features=MAX_FEATURES_PER_ANALYZER,
                ),
            ),
            (
                "char",
                TfidfVectorizer(
                    analyzer="char_wb",
                    ngram_range=CHAR_NGRAM_RANGE,
                    sublinear_tf=True,
                    min_df=2,
                    max_features=MAX_FEATURES_PER_ANALYZER,
                ),
            ),
        ]
    )
    classifier = LogisticRegression(
        C=5.0,
        class_weight="balanced",
        max_iter=MAX_ITERATIONS,
    )
    return Pipeline([("features", features), ("classifier", classifier)])


def build_probability_model():
    return CalibratedClassifierCV(
        estimator=build_base_model(),
        method="sigmoid",
        cv=CALIBRATION_FOLDS,
    )


class DarkPatternModelService:
    def __init__(self, dataset_path):
        self.dataset = self._load_dataset(dataset_path)
        self.metrics = self._evaluate()

        self.explanation_model = build_base_model()
        self.explanation_model.fit(
            self.dataset["text"], self.dataset["Pattern Category"]
        )

        self.probability_model = build_probability_model()
        self.probability_model.fit(
            self.dataset["text"], self.dataset["Pattern Category"]
        )
        self.classes = list(self.probability_model.classes_)

    @staticmethod
    def _load_dataset(dataset_path):
        dataset = pd.read_csv(dataset_path)
        required_columns = {"page_id", "text", "Pattern Category"}
        missing = required_columns.difference(dataset.columns)
        if missing:
            raise ValueError(f"Dataset is missing required columns: {sorted(missing)}")

        dataset = dataset.dropna(subset=list(required_columns)).copy()
        dataset["text"] = dataset["text"].astype(str).str.strip()
        dataset["Pattern Category"] = (
            dataset["Pattern Category"].astype(str).str.strip()
        )
        dataset["page_id"] = dataset["page_id"].astype(str)
        return dataset[dataset["text"].str.len() >= 3].reset_index(drop=True)

    def _evaluate(self):
        splitter = StratifiedGroupKFold(
            n_splits=EVALUATION_FOLDS,
            shuffle=True,
            random_state=RANDOM_STATE,
        )

        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message="The least populated class in y has only",
                category=UserWarning,
            )
            train_index, test_index = next(
                splitter.split(
                    self.dataset["text"],
                    self.dataset["Pattern Category"],
                    groups=self.dataset["page_id"],
                )
            )

        train_data = self.dataset.iloc[train_index]
        test_data = self.dataset.iloc[test_index]
        train_groups = set(train_data["page_id"])
        test_groups = set(test_data["page_id"])
        group_overlap = train_groups.intersection(test_groups)
        if group_overlap:
            raise RuntimeError("Grouped evaluation contains train/test page overlap")

        evaluation_model = build_probability_model()
        evaluation_model.fit(
            train_data["text"], train_data["Pattern Category"]
        )
        predictions = evaluation_model.predict(test_data["text"])
        probabilities = evaluation_model.predict_proba(test_data["text"])

        report = classification_report(
            test_data["Pattern Category"],
            predictions,
            output_dict=True,
            zero_division=0,
        )
        labels = sorted(self.dataset["Pattern Category"].unique())
        matrix = confusion_matrix(
            test_data["Pattern Category"], predictions, labels=labels
        )

        feature_probe = build_base_model()
        feature_probe.fit(train_data["text"], train_data["Pattern Category"])
        feature_names = feature_probe.named_steps["features"].get_feature_names_out()
        vocabulary_size = len(feature_names)
        word_feature_count = sum(
            name.startswith("word__") for name in feature_names
        )
        char_feature_count = vocabulary_size - word_feature_count

        class_distribution = (
            self.dataset["Pattern Category"].value_counts().to_dict()
        )
        rare_classes = {
            label: int(count)
            for label, count in class_distribution.items()
            if count < MIN_RECOMMENDED_CLASS_SIZE
        }

        return {
            "modelName": (
                "Word TF-IDF (1-3 grams) + character TF-IDF (3-5 grams) "
                "+ calibrated Logistic Regression"
            ),
            "datasetSize": int(len(self.dataset)),
            "numClasses": int(self.dataset["Pattern Category"].nunique()),
            "trainSize": int(len(train_data)),
            "testSize": int(len(test_data)),
            "testFraction": round(len(test_data) / len(self.dataset), 4),
            "trainGroupCount": int(len(train_groups)),
            "testGroupCount": int(len(test_groups)),
            "groupOverlap": 0,
            "splitMethod": (
                "StratifiedGroupKFold held-out fold grouped by page_id"
            ),
            "accuracy": round(accuracy_score(test_data["Pattern Category"], predictions), 4),
            "macroF1": round(report["macro avg"]["f1-score"], 4),
            "weightedF1": round(report["weighted avg"]["f1-score"], 4),
            "weightedPrecision": round(report["weighted avg"]["precision"], 4),
            "weightedRecall": round(report["weighted avg"]["recall"], 4),
            "logLoss": round(
                log_loss(
                    test_data["Pattern Category"],
                    probabilities,
                    labels=list(evaluation_model.classes_),
                ),
                4,
            ),
            "calibrated": True,
            "calibrationMethod": (
                f"sigmoid calibration with {CALIBRATION_FOLDS}-fold CV"
            ),
            "vocabularySize": int(vocabulary_size),
            "wordFeatureCount": int(word_feature_count),
            "charFeatureCount": int(char_feature_count),
            "ngramRange": list(WORD_NGRAM_RANGE),
            "charNgramRange": list(CHAR_NGRAM_RANGE),
            "maxIterations": MAX_ITERATIONS,
            "perClass": {
                label: {
                    "precision": round(stats["precision"], 3),
                    "recall": round(stats["recall"], 3),
                    "f1": round(stats["f1-score"], 3),
                    "support": int(stats["support"]),
                }
                for label, stats in report.items()
                if label not in ("accuracy", "macro avg", "weighted avg")
            },
            "classDistribution": {
                label: int(count) for label, count in class_distribution.items()
            },
            "rareClasses": rare_classes,
            "minimumRecommendedClassSize": MIN_RECOMMENDED_CLASS_SIZE,
            "confusionMatrix": {
                "labels": labels,
                "matrix": matrix.tolist(),
            },
        }

    def classify(self, text):
        return self.classify_many([text])[0]

    def classify_many(self, texts):
        cleaned = [" ".join(str(text).split()) for text in texts]
        if not cleaned:
            return []

        predictions = self.probability_model.predict(cleaned)
        probabilities = self.probability_model.predict_proba(cleaned)
        results = []

        for text, prediction, class_probabilities in zip(
            cleaned, predictions, probabilities
        ):
            probability_by_class = dict(zip(self.classes, class_probabilities))
            confidence = float(probability_by_class[prediction])
            top_classes = sorted(
                probability_by_class.items(),
                key=lambda item: item[1],
                reverse=True,
            )[:3]

            results.append(
                {
                    "text": text,
                    "prediction": prediction,
                    "confidence": round(confidence * 100, 1),
                    "confidenceBand": self._confidence_band(confidence),
                    "isDarkPattern": prediction != "Not Dark Pattern",
                    "topClasses": [
                        {
                            "label": label,
                            "probability": round(float(probability) * 100, 1),
                        }
                        for label, probability in top_classes
                    ],
                    "explanation": self.explain(text, prediction),
                    "calibrated": True,
                }
            )

        return results

    def explain(self, text, predicted_class, top_n=5):
        try:
            features = self.explanation_model.named_steps["features"]
            classifier = self.explanation_model.named_steps["classifier"]
            class_index = list(classifier.classes_).index(predicted_class)
            vector = features.transform([text])
            feature_names = features.get_feature_names_out()
            contributions = []

            for column, value in zip(vector.indices, vector.data):
                feature_name = feature_names[column]
                if not feature_name.startswith("word__"):
                    continue
                contribution = float(
                    value * classifier.coef_[class_index][column]
                )
                if contribution > 0:
                    contributions.append(
                        (feature_name.removeprefix("word__"), contribution)
                    )

            contributions.sort(key=lambda item: item[1], reverse=True)
            return [
                {"phrase": phrase, "weight": round(weight, 4)}
                for phrase, weight in contributions[:top_n]
            ]
        except (ValueError, AttributeError):
            return []

    @staticmethod
    def _confidence_band(confidence):
        if confidence >= 0.75:
            return "high"
        if confidence >= 0.50:
            return "moderate"
        return "low"