()
| 66 | |
| 67 | |
| 68 | def main(): |
| 69 | d = Path("uncommon_route/data/v2_splits") |
| 70 | train_texts, train_labels = load_split(d / "train.jsonl") |
| 71 | holdout_texts, holdout_labels = load_split(d / "holdout.jsonl") |
| 72 | tier_names = {0: "LOW", 1: "MID", 2: "MID_HIGH", 3: "HIGH"} |
| 73 | |
| 74 | print(f"Train: {len(train_texts)}, dist={Counter(train_labels)}") |
| 75 | print(f"Holdout: {len(holdout_texts)}, dist={Counter(holdout_labels)}") |
| 76 | |
| 77 | # ─── Baseline ─── |
| 78 | from sentence_transformers import SentenceTransformer, InputExample, losses |
| 79 | from torch.utils.data import DataLoader |
| 80 | |
| 81 | print("\n=== Baseline: Frozen bge-small + LogReg ===") |
| 82 | model = SentenceTransformer("BAAI/bge-small-en-v1.5") |
| 83 | train_emb = model.encode(train_texts, normalize_embeddings=True, show_progress_bar=False) |
| 84 | holdout_emb = model.encode(holdout_texts, normalize_embeddings=True, show_progress_bar=False) |
| 85 | |
| 86 | clf = LogisticRegressionCV(max_iter=2000, random_state=42) |
| 87 | clf.fit(train_emb, train_labels) |
| 88 | base_preds = clf.predict(holdout_emb) |
| 89 | base_acc = accuracy_score(holdout_labels, base_preds) |
| 90 | print(f" Accuracy: {base_acc:.1%}") |
| 91 | for t in range(4): |
| 92 | mask = [l == t for l in holdout_labels] |
| 93 | if any(mask): |
| 94 | ta = accuracy_score([l for l, m in zip(holdout_labels, mask) if m], |
| 95 | [p for p, m in zip(base_preds, mask) if m]) |
| 96 | print(f" {tier_names[t]:8s} n={sum(mask):3d} acc={ta:.0%}") |
| 97 | |
| 98 | # ─── Fine-tune with contrastive learning ─── |
| 99 | print("\n=== Fine-tuning bge-small with contrastive pairs ===") |
| 100 | pairs = make_contrastive_pairs(train_texts, train_labels, n_pairs=8000, seed=42) |
| 101 | print(f" Generated {len(pairs)} contrastive pairs") |
| 102 | |
| 103 | train_examples = [InputExample(texts=[a, b], label=l) for a, b, l in pairs] |
| 104 | train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=32) |
| 105 | train_loss = losses.CosineSimilarityLoss(model) |
| 106 | |
| 107 | print(" Training (3 epochs)...") |
| 108 | model.fit( |
| 109 | train_objectives=[(train_dataloader, train_loss)], |
| 110 | epochs=3, |
| 111 | warmup_steps=100, |
| 112 | show_progress_bar=True, |
| 113 | ) |
| 114 | |
| 115 | # ─── Evaluate fine-tuned ─── |
| 116 | print("\n=== Fine-tuned bge-small + LogReg ===") |
| 117 | ft_train_emb = model.encode(train_texts, normalize_embeddings=True, show_progress_bar=False) |
| 118 | ft_holdout_emb = model.encode(holdout_texts, normalize_embeddings=True, show_progress_bar=False) |
| 119 | |
| 120 | clf_ft = LogisticRegressionCV(max_iter=2000, random_state=42) |
| 121 | clf_ft.fit(ft_train_emb, train_labels) |
| 122 | ft_preds = clf_ft.predict(ft_holdout_emb) |
| 123 | ft_acc = accuracy_score(holdout_labels, ft_preds) |
| 124 | print(f" Accuracy: {ft_acc:.1%}") |
| 125 | for t in range(4): |
no test coverage detected