"""Small, reproducible binary-classification comparison. Run: python benchmark.py Run with Jev: TYPESAFE_API_KEY=... python benchmark.py --jev Jev gets eight labeled training examples from each class, whereas the sklearn models fit the full training split. Treat the results as a practical comparison, not a controlled comparison of learning algorithms. """ import argparse from concurrent.futures import ThreadPoolExecutor import json import os from pathlib import Path import time from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen import numpy as np from sklearn.datasets import load_breast_cancer, load_iris, load_wine from sklearn.ensemble import ExtraTreesClassifier, GradientBoostingClassifier from sklearn.metrics import accuracy_score, confusion_matrix, recall_score, roc_auc_score from sklearn.model_selection import train_test_split SEED = 18 TEST_SIZE = 0.25 EXAMPLES_PER_CLASS = 8 API_URL = "https://api.typesafe.ai/v1/systemone" def datasets(): iris = load_iris() iris_mask = iris.target != 0 cancer = load_breast_cancer() wine = load_wine() return [ ("iris_virginica_vs_versicolor", iris.data[iris_mask], (iris.target[iris_mask] == 2).astype(int), iris.feature_names, "virginica iris", "versicolor iris"), ("breast_cancer_malignant", cancer.data, (cancer.target == 0).astype(int), cancer.feature_names, "malignant sample", "benign sample"), ("wine_class_0", wine.data, (wine.target == 0).astype(int), wine.feature_names, "dataset wine class 0", "dataset wine class 1 or 2"), ] def metrics(y_true, probability): predicted = (np.asarray(probability) >= 0.5).astype(int) tn, fp, fn, tp = confusion_matrix(y_true, predicted, labels=[0, 1]).ravel() return { "accuracy": round(float(accuracy_score(y_true, predicted)), 4), "auc": round(float(roc_auc_score(y_true, probability)), 4), "sensitivity": round(float(recall_score(y_true, predicted)), 4), "specificity": round(float(tn / (tn + fp)), 4), "confusion_matrix": [[int(tn), int(fp)], [int(fn), int(tp)]], } def row_dict(row, feature_names): return {str(name): round(float(value), 4) for name, value in zip(feature_names, row)} def jev_probability(sample, feature_names, examples, positive, negative, api_key): state = { "task": f"Classify the sample as {positive} (positive) or {negative} (negative).", "labeled_training_examples": examples, "sample": row_dict(sample, feature_names), } payload = { "model": "jev-latest", "state": state, "questions": { "positive": { "type": "noul", "instructions": ( f"Based on the labeled examples and the sample's numeric features, " f"is the sample in the positive class ({positive}) rather than " f"the negative class ({negative})?" ), } }, } request = Request( API_URL, data=json.dumps(payload).encode(), headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, method="POST", ) for attempt in range(5): try: with urlopen(request, timeout=60) as response: result = json.load(response) return float(result["answers"]["positive"]["noul"]) except HTTPError as exc: if exc.code not in (429, 500, 502, 503, 504) or attempt == 4: raise RuntimeError(f"Jev API returned HTTP {exc.code}") from exc except URLError: if attempt == 4: raise time.sleep(2 ** attempt) raise AssertionError("unreachable") def read_api_key(env_file): if env_file: for line in env_file.read_text().splitlines(): if "=" not in line or line.lstrip().startswith("#"): continue name, value = line.split("=", 1) if name.strip().removeprefix("export ").lower() in ("typesafe_api_key", "jev_api_key"): return value.strip().strip('"\'') return os.getenv("TYPESAFE_API_KEY") def benchmark(include_jev, env_file=None): api_key = read_api_key(env_file) if include_jev and not api_key: raise SystemExit("--jev requires TYPESAFE_API_KEY in the environment") report = {} for name, X, y, feature_names, positive, negative in datasets(): X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=TEST_SIZE, stratify=y, random_state=SEED ) rng = np.random.default_rng(SEED) chosen = np.concatenate([ rng.choice(np.flatnonzero(y_train == label), EXAMPLES_PER_CLASS, replace=False) for label in (0, 1) ]) entries = {} models = { "gradient_boosting_full": (GradientBoostingClassifier(random_state=SEED), None), "extra_trees_full": (ExtraTreesClassifier(n_estimators=200, random_state=SEED, n_jobs=-1), None), "gradient_boosting_16_shot": (GradientBoostingClassifier(random_state=SEED), chosen), "extra_trees_16_shot": (ExtraTreesClassifier(n_estimators=200, random_state=SEED, n_jobs=-1), chosen), } for model_name, (model, subset) in models.items(): model.fit(X_train if subset is None else X_train[subset], y_train if subset is None else y_train[subset]) entries[model_name] = metrics(y_test, model.predict_proba(X_test)[:, 1]) if include_jev: examples = [ {"features": row_dict(X_train[i], feature_names), "label": "positive" if y_train[i] else "negative"} for i in chosen ] print(f"Running Jev on {name}: {len(X_test)} test rows", flush=True) with ThreadPoolExecutor(max_workers=5) as executor: probabilities = list(executor.map( lambda row: jev_probability( row, feature_names, examples, positive, negative, api_key ), X_test, )) entries["jev_few_shot"] = metrics(y_test, probabilities) print(f"Completed Jev on {name}", flush=True) report[name] = { "n_total": len(y), "n_train": len(y_train), "n_test": len(y_test), "n_few_shot": len(chosen), "n_test_positive": int(y_test.sum()), "positive": positive, "negative": negative, "models": entries, } return report def main(): parser = argparse.ArgumentParser() parser.add_argument("--jev", action="store_true", help="also query Jev on each test row") parser.add_argument("--env-file", type=Path, help="read only jev_api_key or TYPESAFE_API_KEY") parser.add_argument("--out", type=Path, default=Path("benchmark_results.json")) args = parser.parse_args() report = benchmark(args.jev, args.env_file) args.out.write_text(json.dumps(report, indent=2) + "\n") print(json.dumps(report, indent=2)) if __name__ == "__main__": main()