| """Evaluate embedding models on held-out Kazakh retrieval + cross-lingual matching. |
| |
| CUDA_VISIBLE_DEVICES=4 venv-emb/bin/python scripts/eval_embed.py |
| |
| Models: v2 (checkpoints/embed-v2/final), base (BAAI/bge-m3), v1 (Darmm/darmm-embedding-multilingual). |
| Writes results/embed_v2.json. |
| """ |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from datasets import load_from_disk |
| from sentence_transformers import SentenceTransformer |
|
|
| MODELS = { |
| "v2": "checkpoints/embed-v2/final", |
| "base-bge-m3": "BAAI/bge-m3", |
| "v1-darmm": "Darmm/darmm-embedding-multilingual", |
| } |
|
|
|
|
| def encode(model, texts): |
| return model.encode(texts, batch_size=128, normalize_embeddings=True, convert_to_tensor=True, show_progress_bar=False) |
|
|
|
|
| def retrieval_metrics(model, queries, docs): |
| q = encode(model, queries) |
| d = encode(model, docs) |
| r1 = r5 = r10 = mrr = 0.0 |
| for i in range(0, len(queries), 512): |
| sims = q[i : i + 512] @ d.T |
| top = sims.topk(10, dim=1).indices.cpu() |
| for row, gold in enumerate(range(i, min(i + 512, len(queries)))): |
| hits = (top[row] == gold).nonzero() |
| if len(hits): |
| rank = int(hits[0][0]) + 1 |
| mrr += 1.0 / rank |
| r10 += 1 |
| r5 += rank <= 5 |
| r1 += rank == 1 |
| n = len(queries) |
| return {"recall@1": round(r1 / n, 4), "recall@5": round(r5 / n, 4), "recall@10": round(r10 / n, 4), "mrr@10": round(mrr / n, 4), "n": n} |
|
|
|
|
| def xl_accuracy(model, kk, other): |
| a = encode(model, kk) |
| b = encode(model, other) |
| pred = (a @ b.T).argmax(dim=1).cpu() |
| acc = (pred == torch.arange(len(kk))).float().mean().item() |
| return {"acc@1": round(acc, 4), "n": len(kk)} |
|
|
|
|
| def main(): |
| ds = load_from_disk("data/kk_pairs") |
| ret = ds["eval_retrieval"] |
| queries, docs = list(ret["query"]), list(ret["doc"]) |
| kkru = ds["eval_xl_kkru"] |
| kken = ds["eval_xl_kken"] |
|
|
| results = {} |
| for name, path in MODELS.items(): |
| model = SentenceTransformer(path, model_kwargs={"torch_dtype": torch.float16} if torch.cuda.is_available() else {}) |
| model.max_seq_length = 192 |
| results[name] = { |
| "kk_retrieval": retrieval_metrics(model, queries, docs), |
| "xl_kk_ru": xl_accuracy(model, list(kkru["kk"]), list(kkru["other"])), |
| "xl_kk_en": xl_accuracy(model, list(kken["kk"]), list(kken["other"])), |
| } |
| print(name, json.dumps(results[name]), flush=True) |
| del model |
| torch.cuda.empty_cache() |
|
|
| Path("results").mkdir(exist_ok=True) |
| Path("results/embed_v2.json").write_text(json.dumps(results, indent=2)) |
| print(json.dumps(results, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
