Q3 v1 standalone training script for Vast

q3_train_v1_vast.py

@@ -0,0 +1,261 @@
+"""Q3 v1 — Contrastive fine-tune of multilingual-e5-large on pre-mined hard negatives.
+
+Run on Vast.ai (RTX PRO 4500 / A100 / 4090). Memory-tight defaults from prior
+notebook OOM debugging: AMP, BATCH_SIZE=16, gradient checkpointing-friendly.
+
+Inputs:
+  - HF: Dharun72/swiss-legal-checkpoints/q3_train_v1.parquet (pre-mined pairs)
+  - Kaggle competition data: laws_de.csv, train.csv (val split), test.csv (encoding)
+
+Outputs (pushed to HF Dharun72/swiss-legal-checkpoints):
+  - laws_emb_v2.npy  (175933 × 1024 float32, L2-normalized)
+  - queries_v2.npz   (test, train, val splits encoded)
+  - q3_v1_eval.json  (recall@40 + novel-cit recall metrics)
+
+Usage on Vast:
+  pip install sentence-transformers==2.7.0 transformers==4.45.2 huggingface-hub kaggle pandas pyarrow
+  export HF_TOKEN=...
+  python _q3_train_v1_vast.py
+"""
+import os, sys, json, re, random, time, gc
+from pathlib import Path
+from collections import defaultdict
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+
+os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True')
+os.environ.setdefault('TOKENIZERS_PARALLELISM', 'false')
+
+from sentence_transformers import SentenceTransformer, InputExample, losses
+from huggingface_hub import hf_hub_download, HfApi
+
+random.seed(42); np.random.seed(42); torch.manual_seed(42)
+if torch.cuda.is_available(): torch.cuda.manual_seed_all(42)
+
+# ---- Config ----
+CFG = {
+    'BASE_MODEL': 'intfloat/multilingual-e5-large',
+    'BATCH_SIZE':     16,    # OOM-safe
+    'ENCODE_BATCH':   32,
+    'EPOCHS':         2,
+    'LEARNING_RATE':  2e-5,
+    'WARMUP_RATIO':   0.1,
+    'N_NEG_PER_POS':  4,     # subsample 4 negs per positive (out of ~21 mined)
+    'HF_REPO':        'Dharun72/swiss-legal-checkpoints',
+    'OUT_DIR':        Path('outputs_q3v1'),
+    'PUSH_TO_HF':     True,
+    'VAL_SPLIT':      0.1,   # hold out 10% of train for validation
+}
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+HF_TOKEN = os.environ.get('HF_TOKEN')
+assert HF_TOKEN, 'HF_TOKEN env required'
+CFG['OUT_DIR'].mkdir(exist_ok=True, parents=True)
+
+print(f'torch {torch.__version__}  cuda {torch.cuda.is_available()}  device={DEVICE}')
+
+# ---- Load data ----
+def kaggle_download(competition='llm-agentic-legal-information-retrieval', target=Path('/tmp/kaggle')):
+    target.mkdir(exist_ok=True, parents=True)
+    if not (target / 'laws_de.csv').exists():
+        os.system(f'kaggle competitions download -c {competition} -p {target} -q && cd {target} && unzip -o -q *.zip')
+    return target
+
+print('\n=== Loading data ===')
+DATA_DIR = Path('/kaggle/input') if Path('/kaggle/input').exists() else kaggle_download()
+print(f'  data dir: {DATA_DIR}')
+
+laws = pd.read_csv(DATA_DIR / 'laws_de.csv').fillna('')
+laws_cits = laws['citation'].astype(str).tolist()
+laws_text = laws['text'].astype(str).tolist()
+print(f'  laws_de: {len(laws_cits):,}')
+
+train = pd.read_csv(DATA_DIR / 'train.csv').fillna('')
+test  = pd.read_csv(DATA_DIR / 'test.csv').fillna('')
+print(f'  train: {len(train)}  test: {len(test)}')
+
+# Pre-mined training data
+print('  Downloading q3_train_v1.parquet from HF...')
+parquet_path = hf_hub_download(repo_id=CFG['HF_REPO'], filename='q3_train_v1.parquet', token=HF_TOKEN)
+df_pairs = pd.read_parquet(parquet_path)
+print(f'  loaded {len(df_pairs):,} (gold,neg-list) rows from {df_pairs["qid"].nunique()} unique queries')
+
+# Hold out 10% of unique qids as val
+all_qids = df_pairs['qid'].unique().tolist()
+random.shuffle(all_qids)
+n_val = int(len(all_qids) * CFG['VAL_SPLIT'])
+val_qids = set(all_qids[:n_val])
+train_qids = set(all_qids[n_val:])
+print(f'  train qids: {len(train_qids)}  val qids: {len(val_qids)}')
+
+# ---- Build training examples ----
+print('\n=== Building training InputExamples ===')
+Q_PFX = 'query: '
+P_PFX = 'passage: '
+
+cit_to_text = dict(zip(laws_cits, laws_text))
+examples = []
+for r in df_pairs.itertuples():
+    if r.qid not in train_qids: continue
+    neg_cits = list(r.hard_neg_cits)
+    if not neg_cits: continue
+    sampled = random.sample(neg_cits, min(CFG['N_NEG_PER_POS'], len(neg_cits)))
+    for nc in sampled:
+        nt = cit_to_text.get(nc, '')
+        if not nt: continue
+        examples.append(InputExample(texts=[Q_PFX + r.query, P_PFX + r.gold_text, P_PFX + nt]))
+
+random.shuffle(examples)
+print(f'  total InputExamples: {len(examples):,}')
+
+# Validation queries + their gold cits (for recall@40 eval)
+val_qid_to_gold = defaultdict(set)
+for r in df_pairs.itertuples():
+    if r.qid in val_qids:
+        val_qid_to_gold[r.qid].add(r.gold_cit)
+val_queries = []
+val_gold_sets = []
+seen_qids = set()
+for r in df_pairs.itertuples():
+    if r.qid in val_qids and r.qid not in seen_qids:
+        seen_qids.add(r.qid)
+        val_queries.append((r.qid, r.query))
+        val_gold_sets.append(val_qid_to_gold[r.qid])
+print(f'  val queries: {len(val_queries)}')
+
+# ---- Load base model ----
+print(f'\n=== Loading {CFG["BASE_MODEL"]} ===')
+model = SentenceTransformer(CFG['BASE_MODEL'], device=DEVICE)
+print(f'  loaded. Max seq len: {model.max_seq_length}')
+
+# Truncate to 256 for speed (queries + cit text usually < 256 tokens)
+model.max_seq_length = 256
+
+# ---- Eval baseline (before training) ----
+def eval_recall_at_k(model, val_queries, val_gold_sets, k=40, chunk=8192):
+    """Manual recall@K — chunked corpus encoding to avoid OOM."""
+    # Encode corpus
+    laws_passages = [P_PFX + t for t in laws_text]
+    print(f'  encoding {len(laws_passages):,} laws passages in chunks of {chunk}...')
+    all_embs = []
+    t0 = time.time()
+    for i in range(0, len(laws_passages), chunk):
+        embs = model.encode(laws_passages[i:i+chunk], batch_size=CFG['ENCODE_BATCH'],
+                            convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
+        all_embs.append(embs)
+        if (i // chunk) % 4 == 0:
+            print(f'    chunk {i:,}/{len(laws_passages):,}  [{time.time()-t0:.0f}s]', flush=True)
+    laws_emb = np.vstack(all_embs).astype(np.float32)
+
+    # Encode queries
+    q_texts = [Q_PFX + q for _, q in val_queries]
+    q_emb = model.encode(q_texts, batch_size=CFG['ENCODE_BATCH'],
+                         convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
+
+    # Cosine sims
+    sims = q_emb @ laws_emb.T
+
+    recalls = []
+    mean_ranks = []
+    for i, (qid, _) in enumerate(val_queries):
+        gold = val_gold_sets[i]
+        if not gold:
+            continue
+        top_k_idx = np.argpartition(-sims[i], k)[:k]
+        top_k_set = set(laws_cits[j] for j in top_k_idx)
+        recall = len(gold & top_k_set) / len(gold)
+        recalls.append(recall)
+
+        # Mean rank of gold cits
+        ranks = []
+        sorted_idx = np.argsort(-sims[i])
+        idx_to_rank = {laws_cits[j]: r for r, j in enumerate(sorted_idx[:1000])}
+        for g in gold:
+            ranks.append(idx_to_rank.get(g, 1000))
+        mean_ranks.append(np.mean(ranks))
+
+    return {
+        'recall@k': np.mean(recalls),
+        'mean_gold_rank': np.mean(mean_ranks),
+        'n_val': len(recalls),
+        'k': k,
+    }, laws_emb
+
+print('\n=== Baseline eval (BEFORE training) ===')
+torch.cuda.empty_cache(); gc.collect()
+baseline, _ = eval_recall_at_k(model, val_queries, val_gold_sets, k=40)
+print(f'  Baseline recall@40 = {baseline["recall@k"]:.4f}  mean_gold_rank = {baseline["mean_gold_rank"]:.1f}')
+
+# Free baseline corpus
+torch.cuda.empty_cache(); gc.collect()
+
+# ---- Training ----
+print('\n=== Training ===')
+loss_fn = losses.MultipleNegativesRankingLoss(model)
+loader = DataLoader(examples, batch_size=CFG['BATCH_SIZE'], shuffle=True, drop_last=True)
+total_steps = len(loader) * CFG['EPOCHS']
+warmup = int(CFG['WARMUP_RATIO'] * total_steps)
+print(f'  steps/epoch: {len(loader)}  total_steps: {total_steps}  warmup: {warmup}')
+
+t0 = time.time()
+model.fit(
+    train_objectives=[(loader, loss_fn)],
+    epochs=CFG['EPOCHS'],
+    warmup_steps=warmup,
+    optimizer_params={'lr': CFG['LEARNING_RATE']},
+    use_amp=True,
+    show_progress_bar=True,
+    output_path=str(CFG['OUT_DIR'] / 'checkpoint'),
+)
+print(f'  Trained in {time.time()-t0:.0f}s')
+
+# ---- Eval after training ----
+print('\n=== Post-training eval ===')
+torch.cuda.empty_cache(); gc.collect()
+post, laws_emb_v2 = eval_recall_at_k(model, val_queries, val_gold_sets, k=40)
+print(f'  Post recall@40 = {post["recall@k"]:.4f}  mean_gold_rank = {post["mean_gold_rank"]:.1f}')
+print(f'  Delta recall@40 = {post["recall@k"] - baseline["recall@k"]:+.4f}')
+
+# ---- Encode test queries + train queries (for downstream use) ----
+print('\n=== Encoding queries (test + train + val) for downstream ===')
+test_q = test['query'].tolist()
+train_q = train['query'].tolist()
+val_q  = [q for _, q in val_queries]
+
+q_test_emb  = model.encode([Q_PFX + q for q in test_q],  batch_size=CFG['ENCODE_BATCH'], normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True)
+q_train_emb = model.encode([Q_PFX + q for q in train_q], batch_size=CFG['ENCODE_BATCH'], normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True)
+q_val_emb   = model.encode([Q_PFX + q for q in val_q],   batch_size=CFG['ENCODE_BATCH'], normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True)
+
+# Save artifacts
+print('\n=== Saving artifacts ===')
+np.save(CFG['OUT_DIR'] / 'laws_emb_v2.npy', laws_emb_v2.astype(np.float32))
+np.savez(CFG['OUT_DIR'] / 'queries_v2.npz',
+         test=q_test_emb.astype(np.float32),
+         train=q_train_emb.astype(np.float32),
+         val=q_val_emb.astype(np.float32),
+         val_qids=np.array([qid for qid, _ in val_queries]))
+eval_data = {'baseline': baseline, 'post': post, 'cfg': {k: v for k, v in CFG.items() if k != 'OUT_DIR'}}
+with open(CFG['OUT_DIR'] / 'q3_v1_eval.json', 'w') as f:
+    json.dump(eval_data, f, indent=2)
+print(f'  saved laws_emb_v2.npy, queries_v2.npz, q3_v1_eval.json to {CFG["OUT_DIR"]}')
+
+# ---- Push to HF ----
+if CFG['PUSH_TO_HF']:
+    print('\n=== Pushing to HF ===')
+    api = HfApi(token=HF_TOKEN)
+    for fn in ['laws_emb_v2.npy', 'queries_v2.npz', 'q3_v1_eval.json']:
+        api.upload_file(
+            path_or_fileobj=str(CFG['OUT_DIR'] / fn),
+            path_in_repo=fn,
+            repo_id=CFG['HF_REPO'],
+            repo_type='model',
+            commit_message=f'Q3 v1 — {fn}',
+        )
+        print(f'  pushed {fn}')
+
+print('\n=== DONE ===')
+print(f'Baseline recall@40: {baseline["recall@k"]:.4f}')
+print(f'Post     recall@40: {post["recall@k"]:.4f}')
+print(f'Delta:               {post["recall@k"] - baseline["recall@k"]:+.4f}')