commitsorter.py

import warnings
import pandas as pd
import joblib

try:
    import cuml
except ImportError:
    pass

try:
    import xgboost
except ImportError:
    pass

try:
    import catboost
except ImportError:
    pass

try:
    import dense_utils
except ImportError:
    pass

warnings.filterwarnings("ignore", module="sklearn")


class CommitClassifier:
    def __init__(self, model_path="model.joblib"):
        self._model_path = model_path
        self._model = None
        self._label_enc = None

    def _ensure(self):
        if self._model is None:
            data = joblib.load(self._model_path)
            self._model = data["model"]
            self._label_enc = data["label_encoder"]
        return self._model, self._label_enc

    def _build_scores(self, probs):
        return sorted(
            zip(self._label_enc.classes_, probs),
            key=lambda x: x[1],
            reverse=True,
        )

    def sort(
        self,
        text,
        files_count=0,
        additions=0,
        deletions=0,
        changed_tests=0,
        changed_docs=0,
        changed_source=0,
        has_tests=False,
        has_docs=False,
        extensions=None,
        directories=None,
    ):
        model, label_enc = self._ensure()
        row = {
            "text": text,
            "files_count": files_count,
            "additions": additions,
            "deletions": deletions,
            "changed_tests": changed_tests,
            "changed_docs": changed_docs,
            "changed_source": changed_source,
            "has_tests": int(has_tests),
            "has_docs": int(has_docs),
            "extensions": " ".join(extensions or []),
            "directories": " ".join(directories or []),
        }
        df = pd.DataFrame([row])
        pred = model.predict(df)[0]
        probs = model.predict_proba(df)[0]
        label = label_enc.inverse_transform([pred])[0]
        return label, self._build_scores(probs)

    def sort_batch(self, records):
        model, label_enc = self._ensure()
        rows = []
        for r in records:
            rows.append({
                "text": r["text"],
                "files_count": r.get("files_count", 0),
                "additions": r.get("additions", 0),
                "deletions": r.get("deletions", 0),
                "changed_tests": r.get("changed_tests", 0),
                "changed_docs": r.get("changed_docs", 0),
                "changed_source": r.get("changed_source", 0),
                "has_tests": int(r.get("has_tests", False)),
                "has_docs": int(r.get("has_docs", False)),
                "extensions": " ".join(r.get("extensions", [])),
                "directories": " ".join(r.get("directories", [])),
            })
        df = pd.DataFrame(rows)
        preds = model.predict(df)
        probs = model.predict_proba(df)

        results = []
        for i in range(len(records)):
            label = label_enc.inverse_transform([preds[i]])[0]
            scores = self._build_scores(probs[i])
            results.append({"label": label, "probs": dict(scores)})
        return results