Generate (text_a, text_b, label) pairs for contrastive learning. label=1.0 if same tier, label=0.0 if different tier.
(texts, labels, n_pairs=5000, seed=42)
| 36 | |
| 37 | |
| 38 | def make_contrastive_pairs(texts, labels, n_pairs=5000, seed=42): |
| 39 | """Generate (text_a, text_b, label) pairs for contrastive learning. |
| 40 | label=1.0 if same tier, label=0.0 if different tier. |
| 41 | """ |
| 42 | rng = random.Random(seed) |
| 43 | by_tier = {} |
| 44 | for i, (t, l) in enumerate(zip(texts, labels)): |
| 45 | by_tier.setdefault(l, []).append(i) |
| 46 | |
| 47 | pairs = [] |
| 48 | tiers = list(by_tier.keys()) |
| 49 | |
| 50 | for _ in range(n_pairs): |
| 51 | if rng.random() < 0.5: |
| 52 | # Positive pair (same tier) |
| 53 | tier = rng.choice(tiers) |
| 54 | if len(by_tier[tier]) < 2: |
| 55 | continue |
| 56 | i, j = rng.sample(by_tier[tier], 2) |
| 57 | pairs.append((texts[i], texts[j], 1.0)) |
| 58 | else: |
| 59 | # Negative pair (different tier) |
| 60 | t1, t2 = rng.sample(tiers, 2) |
| 61 | i = rng.choice(by_tier[t1]) |
| 62 | j = rng.choice(by_tier[t2]) |
| 63 | pairs.append((texts[i], texts[j], 0.0)) |
| 64 | |
| 65 | return pairs |
| 66 | |
| 67 | |
| 68 | def main(): |