()
| 40 | |
| 41 | |
| 42 | def main(): |
| 43 | d = Path("uncommon_route/data/v2_splits") |
| 44 | train_texts, train_labels, _ = load_split(d / "train.jsonl") |
| 45 | holdout_texts, holdout_labels, holdout_benchmarks = load_split(d / "holdout.jsonl") |
| 46 | |
| 47 | print(f"Train: {len(train_texts)} samples, dist={Counter(train_labels)}") |
| 48 | print(f"Holdout: {len(holdout_texts)} samples, dist={Counter(holdout_labels)}") |
| 49 | |
| 50 | tier_names = {0: "LOW", 1: "MID", 2: "MID_HIGH", 3: "HIGH"} |
| 51 | |
| 52 | # ─── Baseline: frozen embeddings + LogReg ─── |
| 53 | print("\n=== Baseline: Frozen bge-small + LogReg ===") |
| 54 | from sentence_transformers import SentenceTransformer |
| 55 | frozen_model = SentenceTransformer("BAAI/bge-small-en-v1.5") |
| 56 | train_emb = frozen_model.encode(train_texts, normalize_embeddings=True, show_progress_bar=False) |
| 57 | holdout_emb = frozen_model.encode(holdout_texts, normalize_embeddings=True, show_progress_bar=False) |
| 58 | |
| 59 | clf_baseline = LogisticRegressionCV(max_iter=2000, random_state=42) |
| 60 | clf_baseline.fit(train_emb, train_labels) |
| 61 | baseline_preds = clf_baseline.predict(holdout_emb) |
| 62 | baseline_acc = accuracy_score(holdout_labels, baseline_preds) |
| 63 | print(f" Overall accuracy: {baseline_acc:.1%}") |
| 64 | for tier in range(4): |
| 65 | mask = [l == tier for l in holdout_labels] |
| 66 | if any(mask): |
| 67 | tier_preds = [p for p, m in zip(baseline_preds, mask) if m] |
| 68 | tier_true = [l for l, m in zip(holdout_labels, mask) if m] |
| 69 | print(f" {tier_names[tier]:8s} n={sum(mask):3d} acc={accuracy_score(tier_true, tier_preds):.0%}") |
| 70 | |
| 71 | # ─── Experiment: SetFit fine-tuned embeddings ─── |
| 72 | print("\n=== SetFit: Fine-tuned bge-small ===") |
| 73 | from setfit import SetFitModel, SetFitTrainer |
| 74 | from datasets import Dataset |
| 75 | |
| 76 | train_ds = Dataset.from_dict({"text": train_texts, "label": train_labels}) |
| 77 | holdout_ds = Dataset.from_dict({"text": holdout_texts, "label": holdout_labels}) |
| 78 | |
| 79 | setfit_model = SetFitModel.from_pretrained( |
| 80 | "BAAI/bge-small-en-v1.5", |
| 81 | labels=list(tier_names.values()), |
| 82 | ) |
| 83 | |
| 84 | trainer = SetFitTrainer( |
| 85 | model=setfit_model, |
| 86 | train_dataset=train_ds, |
| 87 | eval_dataset=holdout_ds, |
| 88 | num_iterations=20, # contrastive pairs per sample |
| 89 | num_epochs=1, |
| 90 | batch_size=16, |
| 91 | seed=42, |
| 92 | ) |
| 93 | |
| 94 | print(" Training...") |
| 95 | trainer.train() |
| 96 | |
| 97 | print(" Evaluating...") |
| 98 | setfit_preds = setfit_model.predict(holdout_texts) |
| 99 | # Convert to list of ints |
no test coverage detected