runner.py

from __future__ import annotations
import contextlib, io, os, re
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)
from br_classification import run as _br_run

import nltk

nltk.download("stopwords", quiet=True)
from nltk.corpus import stopwords

_HTML_RE = re.compile(r"<.*?>")
_EMOJI_RE = re.compile(
    "["
    "\U0001f600-\U0001f64f"
    "\U0001f300-\U0001f5ff"
    "\U0001f680-\U0001f6ff"
    "\U0001f1e0-\U0001f1ff"
    "\U00002702-\U000027b0"
    "\U000024c2-\U0001f251"
    "]+",
    flags=re.UNICODE,
)
_KEEP = {
    "no",
    "not",
    "nor",
    "never",
    "very",
    "more",
    "most",
    "too",
    "up",
    "down",
    "off",
    "over",
    "under",
    "against",
    "why",
    "how",
    "when",
}
_STOP_IMPROVED = set(stopwords.words("english")) - _KEEP
_STOP_BASELINE = set(stopwords.words("english")) | {"..."}

ALL_PROJECTS = ["pytorch", "tensorflow", "keras", "incubator-mxnet", "caffe"]


def clean_baseline(text: str) -> str:
    text = _HTML_RE.sub("", str(text))
    text = _EMOJI_RE.sub("", text)
    text = " ".join(w for w in text.split() if w not in _STOP_BASELINE)
    text = re.sub(r"[^A-Za-z0-9(),.!?'`]", " ", text)
    text = re.sub(r"\'s", " 's", text)
    text = re.sub(r"\'ve", " 've", text)
    text = re.sub(r"\)", " ) ", text)
    text = re.sub(r"\?", " ? ", text)
    text = re.sub(r"\s{2,}", " ", text)
    text = re.sub(r"\\", "", text)
    text = re.sub(r"\'", "", text)
    text = re.sub(r"\"", "", text)
    return text.strip().lower()


def clean_improved(text: str) -> str:
    text = _HTML_RE.sub(" ", str(text))
    text = _EMOJI_RE.sub(" ", text)
    text = re.sub(r"([a-z])([A-Z])", r"\1 \2", text)
    text = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1 \2", text)
    text = re.sub(r"[^A-Za-z0-9%]+", " ", text).lower()
    text = re.sub(r"\s+", " ", text).strip()
    return " ".join(w for w in text.split() if w not in _STOP_IMPROVED and len(w) > 1)


def load_project(project: str) -> pd.DataFrame:
    root = os.path.dirname(os.path.abspath(__file__))
    df = pd.read_csv(os.path.join(root, "datasets", f"{project}.csv")).fillna("")
    df["text"] = df.apply(
        lambda r: (
            str(r["Title"]) + ". " + str(r["Body"]) if r["Body"] else str(r["Title"])
        ),
        axis=1,
    )
    df["title"] = df["Title"].astype(str)
    return df[["text", "title", "class"]].rename(columns={"class": "label"})


def run_br_classification(project: str, repeats: int) -> dict:
    with contextlib.redirect_stdout(io.StringIO()):
        return _br_run(project, repeats)


METRICS = ["Accuracy", "Precision", "Recall", "F1", "AUC"]


def _compute(y_true, y_pred, y_prob) -> dict:
    two = len(np.unique(y_true)) > 1
    return {
        "Accuracy": accuracy_score(y_true, y_pred),
        "Precision": precision_score(y_true, y_pred, average="macro", zero_division=0),
        "Recall": recall_score(y_true, y_pred, average="macro", zero_division=0),
        "F1": f1_score(y_true, y_pred, average="macro", zero_division=0),
        "AUC": roc_auc_score(y_true, y_prob) if two else float("nan"),
    }


def run_model(fn, prep: str, project: str, repeats: int) -> dict:
    df = load_project(project)
    clean = clean_baseline if prep == "baseline" else clean_improved
    text = df["text"].apply(clean)
    title = df["title"].apply(clean)
    runs = []
    for seed in range(repeats):
        tr, te = train_test_split(
            np.arange(len(df)), test_size=0.3, random_state=seed, stratify=df["label"]
        )
        y_pred, y_prob = fn(
            text.iloc[tr],
            df["label"].iloc[tr].values,
            text.iloc[te],
            seed,
            tr_title=title.iloc[tr],
            te_title=title.iloc[te],
        )
        runs.append(_compute(df["label"].iloc[te].values, y_pred, y_prob))
    result = {k: float(np.mean([r[k] for r in runs])) for k in METRICS}
    result["_auc_list"] = [r["AUC"] for r in runs]
    return result


def print_table(results: dict[str, dict], project: str):
    names = list(results.keys())
    baseline = names[0]
    col_w = 12

    print(f"\n  Project: {project}")
    header = f"  {'Metric':<11}" + "".join(f"{n:>{col_w}}" for n in names)
    if len(names) > 1:
        header += "".join(f"{'Δ '+n:>{col_w}}" for n in names[1:])
    print(header)
    print(f"  {'-' * (11 + col_w * len(names) + col_w * (len(names)-1))}")

    for k in METRICS:
        row = f"  {k:<11}"
        for n in names:
            row += f"{results[n][k]:>{col_w}.4f}"
        for n in names[1:]:
            delta = results[n][k] - results[baseline][k]
            row += f"{delta:>+{col_w}.4f}"
        print(row)