"""LoRA vs full fine-tuning of a small embedding model on SciFact (BEIR).

Base: BAAI/bge-small-en-v1.5. Train on the 919 train qrels (query, positive
abstract) with MultipleNegativesRankingLoss; evaluate on the 300 held-out test
queries over the full 5,183-abstract corpus. Fixed seed.
Usage: python3 run.py <arm>   arm in: base, lora_qv, lora_all, lora_all_hi, full
Published with the results at https://mixpeek.com/measurements/2026-11-04-lora-embedding-finetune.json
Change device="mps" to "cuda" or "cpu" for your machine.
"""
import json, math, os, sys, time, shutil
import numpy as np
import torch
from datasets import load_dataset, Dataset
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments, losses

torch.set_num_threads(4)
torch.manual_seed(0)
ARM = sys.argv[1]
HERE = os.path.dirname(os.path.abspath(__file__))
MODEL = "BAAI/bge-small-en-v1.5"
QP = "Represent this sentence for searching relevant passages: "

corpus = load_dataset("mteb/scifact", "corpus", split="corpus")
queries = {r["_id"]: r["text"] for r in load_dataset("mteb/scifact", "queries", split="queries")}
docs = {r["_id"]: (r["title"] + ". " + r["text"]).strip() for r in corpus}
doc_ids = list(docs)
def qrels(split):
    out = {}
    for r in load_dataset("mteb/scifact", "default", split=split):
        if r["score"] > 0:
            out.setdefault(str(r["query-id"]), set()).add(str(r["corpus-id"]))
    return out
train, test = qrels("train"), qrels("test")
assert not set(train) & set(test)

def evaluate(model):
    t = time.time()
    D = model.encode([docs[d] for d in doc_ids], batch_size=64, normalize_embeddings=True, convert_to_numpy=True)
    qids = list(test)
    Q = model.encode([QP + queries[q] for q in qids], batch_size=64, normalize_embeddings=True, convert_to_numpy=True)
    S = Q @ D.T
    ndcg, rec = [], []
    for i, q in enumerate(qids):
        top = np.argsort(-S[i])[:10]
        rel = test[q]
        gains = [1.0 if doc_ids[j] in rel else 0.0 for j in top]
        dcg = sum(g / math.log2(k + 2) for k, g in enumerate(gains))
        idcg = sum(1 / math.log2(k + 2) for k in range(min(len(rel), 10)))
        ndcg.append(dcg / idcg)
        rec.append(sum(gains) / len(rel))
    return {"ndcg@10": round(float(np.mean(ndcg)), 4), "recall@10": round(float(np.mean(rec)), 4),
            "queries": len(qids), "docs": len(doc_ids), "eval_s": round(time.time() - t, 1)}

model = SentenceTransformer(MODEL, device="mps")
total = sum(p.numel() for p in model.parameters())
result = {"arm": ARM, "model": MODEL, "total_params": total}

if ARM != "base":
    if ARM.startswith("lora"):
        from peft import LoraConfig, TaskType
        target = ["query", "value"] if ARM.startswith("lora_qv") else "all-linear"
        model.add_adapter(LoraConfig(task_type=TaskType.FEATURE_EXTRACTION, r=16, lora_alpha=16,
                                     lora_dropout=0.1, target_modules=target))
        lr = 1e-3 if ARM.endswith("_hi") else 1e-4
    else:
        lr = 2e-5
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    pairs = [(QP + queries[q], docs[d]) for q, ds in train.items() for d in ds]
    ds = Dataset.from_dict({"anchor": [a for a, _ in pairs], "positive": [p for _, p in pairs]})
    out = os.path.join(HERE, "out", ARM)
    args = SentenceTransformerTrainingArguments(
        output_dir=out, num_train_epochs=3, per_device_train_batch_size=32, learning_rate=lr,
        warmup_ratio=0.1, seed=0, logging_steps=50, save_strategy="no", report_to=[],
        batch_sampler="no_duplicates")
    t = time.time()
    SentenceTransformerTrainer(model=model, args=args, train_dataset=ds,
                               loss=losses.MultipleNegativesRankingLoss(model)).train()
    result.update(train_s=round(time.time() - t, 1), train_pairs=len(pairs), lr=lr, epochs=3,
                  trainable_params=trainable, trainable_pct=round(100 * trainable / total, 2))
    save = os.path.join(HERE, "saved", ARM)
    shutil.rmtree(save, ignore_errors=True)
    model.save_pretrained(save)
    size = 0
    for root, _, files in os.walk(save):
        for f in files:
            if f.endswith((".safetensors", ".bin")):
                size += os.path.getsize(os.path.join(root, f))
    result["weights_on_disk_mb"] = round(size / 1e6, 2)

result.update(evaluate(model))
print(json.dumps(result))
with open(os.path.join(HERE, "results.jsonl"), "a") as f:
    f.write(json.dumps(result) + "\n")