train_compare.py

"""
Train and compare emotion classifiers on original vs. SDVM-refined data.

Follows the NLP with Transformers book (Chapter 2) approach:
  - Emotion classification on dair-ai/emotion-style data
  - Compare TF-IDF + Logistic Regression trained on original vs. SDVM-refined text
  - Metrics: accuracy, macro F1, log-loss (cross-entropy), per-class F1

Requirements:
  pip install sdvm scikit-learn

Environment:
  export SDVM_API_KEY="your-key-here"

Pipeline:
  1. Generate 120 labeled emotion samples (20/class) or load from cache
  2. Refine 90 training samples using SDVM API (cached)
  3. Train TF-IDF + LR on original and refined training sets
  4. Evaluate both on same 30-sample test set (unrefined)
  5. Save results to train_results.json
"""

import json
import math
import os
import time
from collections import defaultdict

from sdvm import Refinery, RawText
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    f1_score,
    log_loss,
)
from sklearn.pipeline import Pipeline

SDVM_API_KEY = os.environ["SDVM_API_KEY"]

EMOTIONS = ["joy", "sadness", "anger", "fear", "surprise", "love"]
SAMPLES_PER_CLASS = 20
TRAIN_PER_CLASS = 15
TEST_PER_CLASS = 5
BATCH_SIZE = 15


def refine_batch(refinery: Refinery, samples: list[str]) -> list[str]:
    """Refine a batch of text samples using SDVM API."""
    results = refinery.run([RawText(text=s) for s in samples])
    return [r.text for r in results]


def train_classifier(texts: list[str], labels: list[str]) -> Pipeline:
    """TF-IDF (1-2 gram) + Logistic Regression -- Ch.2 NLP with Transformers baseline."""
    pipe = Pipeline([
        ("tfidf", TfidfVectorizer(ngram_range=(1, 2), min_df=1, max_features=10000)),
        ("lr", LogisticRegression(max_iter=1000, C=1.0, solver="lbfgs")),
    ])
    pipe.fit(texts, labels)
    return pipe


def evaluate_classifier(pipe: Pipeline, texts: list[str], labels: list[str], name: str) -> dict:
    preds = pipe.predict(texts)
    probs = pipe.predict_proba(texts)
    acc = accuracy_score(labels, preds)
    macro_f1 = f1_score(labels, preds, average="macro")
    ll = log_loss(labels, probs)
    report = classification_report(labels, preds, output_dict=True)
    per_class_f1 = {
        e: round(report.get(e, {}).get("f1-score", 0.0), 4)
        for e in EMOTIONS
    }
    print(f"\n{'='*50}")
    print(f"Results: {name}")
    print(f"  Accuracy:  {acc:.4f}")
    print(f"  Macro F1:  {macro_f1:.4f}")
    print(f"  Log-loss:  {ll:.4f}")
    print(f"  Per-class F1: {per_class_f1}")
    return {
        "accuracy": round(acc, 4),
        "macro_f1": round(macro_f1, 4),
        "log_loss": round(ll, 4),
        "per_class_f1": per_class_f1,
    }


def main():
    refinery = Refinery(api_key=SDVM_API_KEY)

    # --- Step 1: Load cached data ---
    generated_path = "train_data_generated.json"
    if os.path.exists(generated_path):
        print("Loading cached generated data...")
        with open(generated_path, encoding="utf-8") as f:
            saved = json.load(f)
        train_texts_orig = [d["text"] for d in saved["train_orig"]]
        train_labels = [d["label"] for d in saved["train_orig"]]
        test_texts = [d["text"] for d in saved["test"]]
        test_labels = [d["label"] for d in saved["test"]]
        print(f"  Train: {len(train_texts_orig)}, Test: {len(test_texts)}")
    else:
        print("ERROR: train_data_generated.json not found.")
        print("Generate labeled samples first or download from the dataset repo.")
        return

    # --- Step 2: Load or refine training data ---
    refined_path = "train_data_refined.json"
    if os.path.exists(refined_path):
        print("\nLoading cached refined data...")
        with open(refined_path, encoding="utf-8") as f:
            train_texts_refined = json.load(f)
        print(f"  Refined: {len(train_texts_refined)} samples")
    else:
        print("\nStep 2: Refining training samples using SDVM API...")
        train_texts_refined = []
        num_batches = math.ceil(len(train_texts_orig) / BATCH_SIZE)
        for i in range(0, len(train_texts_orig), BATCH_SIZE):
            batch = train_texts_orig[i:i + BATCH_SIZE]
            batch_num = i // BATCH_SIZE + 1
            print(f"  Refining batch {batch_num}/{num_batches} ({len(batch)} samples)...")
            refined = refine_batch(refinery, batch)
            train_texts_refined.extend(refined)
            time.sleep(0.5)
        with open(refined_path, "w", encoding="utf-8") as f:
            json.dump(train_texts_refined, f, indent=2, ensure_ascii=False)
        print(f"  Refined {len(train_texts_refined)} samples -- saved to {refined_path}")

    # --- Step 3: Train classifiers ---
    print("\nStep 3: Training classifiers (TF-IDF + Logistic Regression)...")
    print("  Training on ORIGINAL data...")
    clf_orig = train_classifier(train_texts_orig, train_labels)
    print("  Training on REFINED data...")
    clf_refined = train_classifier(train_texts_refined, train_labels)

    # --- Step 4: Evaluate ---
    print("\nStep 4: Evaluating both classifiers on test set...")
    metrics_orig = evaluate_classifier(clf_orig, test_texts, test_labels, "Original training data")
    metrics_refined = evaluate_classifier(clf_refined, test_texts, test_labels, "Refined (SDVM) training data")

    # --- Build results ---
    improvement = {
        "accuracy_delta": round(metrics_refined["accuracy"] - metrics_orig["accuracy"], 4),
        "macro_f1_delta": round(metrics_refined["macro_f1"] - metrics_orig["macro_f1"], 4),
    }

    results = {
        "experiment": {
            "task": "Emotion classification (dair-ai/emotion style)",
            "model": "TF-IDF (1-2 gram, max 10K features) + Logistic Regression",
            "reference": "NLP with Transformers Ch. 2 -- Text Classification baseline",
            "train_samples": len(train_texts_orig),
            "test_samples": len(test_texts),
            "classes": EMOTIONS,
        },
        "original_training": {"metrics": metrics_orig},
        "refined_training": {"metrics": metrics_refined},
        "improvement": improvement,
    }

    with open("train_results.json", "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    print("\n" + "=" * 60)
    print("FINAL COMPARISON")
    print(f"  Accuracy:   {metrics_orig['accuracy']:.4f} -> {metrics_refined['accuracy']:.4f} ({improvement['accuracy_delta']:+.4f})")
    print(f"  Macro F1:   {metrics_orig['macro_f1']:.4f} -> {metrics_refined['macro_f1']:.4f} ({improvement['macro_f1_delta']:+.4f})")
    print("\nResults written to train_results.json")


if __name__ == "__main__":
    main()