""" 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()