Spaces:

mostafaelsaeed
/

clause-ai

Running

clause-ai / scripts /run_classical_baseline.py

experiment: add classical TF-IDF baseline script

68f2a2d 2 months ago

6.53 kB

#!/usr/bin/env python
"""
Stage 12: Classical (non-LLM) baseline for comparison.
"Do not assume an LLM is automatically superior" -- this fits a TF-IDF +
Logistic Regression multi-label classifier on the REAL training split
(356 contracts) and evaluates on the REAL test split (77 contracts) using
the exact same (contract, category) presence protocol and evaluator as the
local-LLM baseline (src/evaluator.py::BaselineEvaluator), so the two
numbers are directly comparable.
Usage:
  python scripts/run_classical_baseline.py
No API, no local model server, no GPU required -- this is a CPU-only
scikit-learn pipeline, useful as a sanity floor: if the LLM baseline can't
beat this, that's a real, reportable finding, not something to hide.
"""

import os
import sys
import json
import logging
import subprocess
from pathlib import Path
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import MultiLabelBinarizer

from src.schemas import CUAD_CATEGORIES
from src.evaluator import BaselineEvaluator

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

RANDOM_SEED = 42  # Same seed used throughout the project (data split, etc.)


def _git_commit_hash() -> str:
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=Path(__file__).resolve().parent.parent,
            stderr=subprocess.DEVNULL,
        ).decode().strip()
    except Exception:
        return "unknown"


def load_split(name: str):
    path = Path(f"data/processed/cuad_{name}_contracts.json")
    if not path.exists():
        raise FileNotFoundError(f"{path} not found -- run scripts/prepare_datasets.py first")
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def contract_labels(contract: dict) -> list:
    """The set of categories actually present in this contract (gold),
    matching the (contract, category) presence semantics used everywhere
    else in this project (ADR-018)."""
    return sorted({c["category"] for c in contract.get("clauses", [])})


def main():
    logger.info("=" * 70)
    logger.info("CLASSICAL BASELINE: TF-IDF + One-vs-Rest Logistic Regression")
    logger.info("=" * 70)

    train_contracts = load_split("train")
    test_contracts = load_split("test")
    logger.info(f"Train: {len(train_contracts)} contracts | Test: {len(test_contracts)} contracts")

    X_train_text = [c["text"] for c in train_contracts]
    X_test_text = [c["text"] for c in test_contracts]
    y_train_labels = [contract_labels(c) for c in train_contracts]
    y_test_labels = [contract_labels(c) for c in test_contracts]  # not used for fitting -- gold for eval only

    # Binarize labels over the FULL 41-category space (not just categories
    # seen in train) so a category absent from train but present in test
    # is handled correctly as "never predicted" rather than silently
    # dropped from the label space.
    mlb = MultiLabelBinarizer(classes=CUAD_CATEGORIES)
    Y_train = mlb.fit_transform(y_train_labels)
    logger.info(f"Label matrix: {Y_train.shape} (contracts x categories)")

    # TF-IDF fit ONLY on train text -- never on test, to avoid any leakage
    # of test vocabulary/IDF statistics into the model.
    vectorizer = TfidfVectorizer(
        max_features=20000,
        ngram_range=(1, 2),
        min_df=2,
        stop_words="english",
    )
    X_train = vectorizer.fit_transform(X_train_text)
    X_test = vectorizer.transform(X_test_text)
    logger.info(f"TF-IDF matrix: train {X_train.shape}, test {X_test.shape}")

    # class_weight='balanced' matters heavily here: CUAD's real category
    # frequencies range from 27 (Price Restrictions... actually 0 in test,
    # but present in train) to 2550 (Parties) instances -- see
    # data/DATA_SOURCES.md. Without balancing, rare categories would never
    # be predicted at all.
    clf = OneVsRestClassifier(
        LogisticRegression(
            max_iter=2000,
            class_weight="balanced",
            random_state=RANDOM_SEED,
        )
    )
    logger.info("Fitting classifier on train split only...")
    clf.fit(X_train, Y_train)

    logger.info("Predicting on test split...")
    Y_pred_proba = clf.predict_proba(X_test)
    Y_pred = clf.predict(X_test)

    # Convert to the same prediction-dict shape the evaluator expects,
    # tagged with contract_id exactly like the LLM pipeline's output, so
    # the identical BaselineEvaluator.compute_f1() call scores both fairly.
    predictions = []
    for i, contract in enumerate(test_contracts):
        for j, category in enumerate(CUAD_CATEGORIES):
            if Y_pred[i, j] == 1:
                predictions.append({
                    "contract_id": contract["id"],
                    "category": category,
                    "confidence": float(Y_pred_proba[i, j]),
                })

    gold = [
        {"contract_id": c["id"], "category": cat}
        for c in test_contracts
        for cat in contract_labels(c)
    ]

    evaluator = BaselineEvaluator()
    metrics = evaluator.compute_f1(predictions, gold)
    calibration = evaluator.compute_confidence_calibration(predictions, gold)
    evaluator.log_results(metrics)

    logger.info(f"\nClassical baseline Macro F1: {metrics['macro_f1']:.4f}")
    logger.info(f"Classical baseline Micro F1: {metrics['micro_f1']:.4f}")

    output = {
        "method": "TF-IDF + One-vs-Rest Logistic Regression",
        "sklearn_version": __import__("sklearn").__version__,
        "random_seed": RANDOM_SEED,
        "git_commit": _git_commit_hash(),
        "timestamp": datetime.utcnow().isoformat(),
        "train_contracts": len(train_contracts),
        "test_contracts": len(test_contracts),
        "tfidf_params": {"max_features": 20000, "ngram_range": [1, 2], "min_df": 2},
        "classifier_params": {"class_weight": "balanced", "max_iter": 2000},
        "metrics": metrics,
        "calibration": calibration,
    }

    out_path = Path("data/processed/classical_baseline_results.json")
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(output, f, indent=2)
    logger.info(f"Results saved to {out_path}")

    return output


if __name__ == "__main__":
    main()

	#!/usr/bin/env python
	"""
	Stage 12: Classical (non-LLM) baseline for comparison.

	"Do not assume an LLM is automatically superior" -- this fits a TF-IDF +
	Logistic Regression multi-label classifier on the REAL training split
	(356 contracts) and evaluates on the REAL test split (77 contracts) using
	the exact same (contract, category) presence protocol and evaluator as the
	local-LLM baseline (src/evaluator.py::BaselineEvaluator), so the two
	numbers are directly comparable.

	Usage:
	python scripts/run_classical_baseline.py

	No API, no local model server, no GPU required -- this is a CPU-only
	scikit-learn pipeline, useful as a sanity floor: if the LLM baseline can't
	beat this, that's a real, reportable finding, not something to hide.
	"""

	import os
	import sys
	import json
	import logging
	import subprocess
	from pathlib import Path
	from datetime import datetime

	sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

	from sklearn.feature_extraction.text import TfidfVectorizer
	from sklearn.linear_model import LogisticRegression
	from sklearn.multiclass import OneVsRestClassifier
	from sklearn.preprocessing import MultiLabelBinarizer

	from src.schemas import CUAD_CATEGORIES
	from src.evaluator import BaselineEvaluator

	logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
	logger = logging.getLogger(__name__)

	RANDOM_SEED = 42 # Same seed used throughout the project (data split, etc.)


	def _git_commit_hash() -> str:
	try:
	return subprocess.check_output(
	["git", "rev-parse", "--short", "HEAD"],
	cwd=Path(__file__).resolve().parent.parent,
	stderr=subprocess.DEVNULL,
	).decode().strip()
	except Exception:
	return "unknown"


	def load_split(name: str):
	path = Path(f"data/processed/cuad_{name}_contracts.json")
	if not path.exists():
	raise FileNotFoundError(f"{path} not found -- run scripts/prepare_datasets.py first")
	with open(path, encoding="utf-8") as f:
	return json.load(f)


	def contract_labels(contract: dict) -> list:
	"""The set of categories actually present in this contract (gold),
	matching the (contract, category) presence semantics used everywhere
	else in this project (ADR-018)."""
	return sorted({c["category"] for c in contract.get("clauses", [])})


	def main():
	logger.info("=" * 70)
	logger.info("CLASSICAL BASELINE: TF-IDF + One-vs-Rest Logistic Regression")
	logger.info("=" * 70)

	train_contracts = load_split("train")
	test_contracts = load_split("test")
	logger.info(f"Train: {len(train_contracts)} contracts \| Test: {len(test_contracts)} contracts")

	X_train_text = [c["text"] for c in train_contracts]
	X_test_text = [c["text"] for c in test_contracts]
	y_train_labels = [contract_labels(c) for c in train_contracts]
	y_test_labels = [contract_labels(c) for c in test_contracts] # not used for fitting -- gold for eval only

	# Binarize labels over the FULL 41-category space (not just categories
	# seen in train) so a category absent from train but present in test
	# is handled correctly as "never predicted" rather than silently
	# dropped from the label space.
	mlb = MultiLabelBinarizer(classes=CUAD_CATEGORIES)
	Y_train = mlb.fit_transform(y_train_labels)
	logger.info(f"Label matrix: {Y_train.shape} (contracts x categories)")

	# TF-IDF fit ONLY on train text -- never on test, to avoid any leakage
	# of test vocabulary/IDF statistics into the model.
	vectorizer = TfidfVectorizer(
	max_features=20000,
	ngram_range=(1, 2),
	min_df=2,
	stop_words="english",
	)
	X_train = vectorizer.fit_transform(X_train_text)
	X_test = vectorizer.transform(X_test_text)
	logger.info(f"TF-IDF matrix: train {X_train.shape}, test {X_test.shape}")

	# class_weight='balanced' matters heavily here: CUAD's real category
	# frequencies range from 27 (Price Restrictions... actually 0 in test,
	# but present in train) to 2550 (Parties) instances -- see
	# data/DATA_SOURCES.md. Without balancing, rare categories would never
	# be predicted at all.
	clf = OneVsRestClassifier(
	LogisticRegression(
	max_iter=2000,
	class_weight="balanced",
	random_state=RANDOM_SEED,
	)
	)
	logger.info("Fitting classifier on train split only...")
	clf.fit(X_train, Y_train)

	logger.info("Predicting on test split...")
	Y_pred_proba = clf.predict_proba(X_test)
	Y_pred = clf.predict(X_test)

	# Convert to the same prediction-dict shape the evaluator expects,
	# tagged with contract_id exactly like the LLM pipeline's output, so
	# the identical BaselineEvaluator.compute_f1() call scores both fairly.
	predictions = []
	for i, contract in enumerate(test_contracts):
	for j, category in enumerate(CUAD_CATEGORIES):
	if Y_pred[i, j] == 1:
	predictions.append({
	"contract_id": contract["id"],
	"category": category,
	"confidence": float(Y_pred_proba[i, j]),
	})

	gold = [
	{"contract_id": c["id"], "category": cat}
	for c in test_contracts
	for cat in contract_labels(c)
	]

	evaluator = BaselineEvaluator()
	metrics = evaluator.compute_f1(predictions, gold)
	calibration = evaluator.compute_confidence_calibration(predictions, gold)
	evaluator.log_results(metrics)

	logger.info(f"\nClassical baseline Macro F1: {metrics['macro_f1']:.4f}")
	logger.info(f"Classical baseline Micro F1: {metrics['micro_f1']:.4f}")

	output = {
	"method": "TF-IDF + One-vs-Rest Logistic Regression",
	"sklearn_version": __import__("sklearn").__version__,
	"random_seed": RANDOM_SEED,
	"git_commit": _git_commit_hash(),
	"timestamp": datetime.utcnow().isoformat(),
	"train_contracts": len(train_contracts),
	"test_contracts": len(test_contracts),
	"tfidf_params": {"max_features": 20000, "ngram_range": [1, 2], "min_df": 2},
	"classifier_params": {"class_weight": "balanced", "max_iter": 2000},
	"metrics": metrics,
	"calibration": calibration,
	}

	out_path = Path("data/processed/classical_baseline_results.json")
	with open(out_path, "w", encoding="utf-8") as f:
	json.dump(output, f, indent=2)
	logger.info(f"Results saved to {out_path}")

	return output


	if __name__ == "__main__":
	main()