| |
| """Reproduce papr-embed-code-1 (developed as papr-embed-v2) MTEB scores from the released package with stock mteb. |
| |
| python replicate_mteb.py --model ./papr-embed-code-1-4b --tasks coir --out results/code-1-4b |
| python replicate_mteb.py --model Qwen/Qwen3-Embedding-4B --prompts task_prompts.json --tasks ArguAna,NFCorpus --out results/base-4b |
| |
| Loads the model the way mteb loads Qwen3-Embedding (q3e_instruct_loader: sentence-transformers, last-token |
| pooling, "Instruct: {instruction}\\nQuery:" on queries only, bare documents) with the package's |
| task_prompts.json. Everything else comes from the package: float32 weights, max_seq_length 8192 |
| (sentence_bert_config.json), tokenizer-default padding. Batch size is mteb's default, 32. A stock Qwen3 |
| baseline is loaded at mteb's pinned revision with the same prompts, fp32 and max_seq_length 8192. |
| """ |
| from __future__ import annotations |
|
|
| import argparse, json, time |
| from pathlib import Path |
|
|
| COIR = ["AppsRetrieval", "COIRCodeSearchNetRetrieval", "CodeFeedbackMT", "CodeFeedbackST", "CodeSearchNetCCRetrieval", |
| "CodeTransOceanContest", "CodeTransOceanDL", "CosQA", "StackOverflowQA", "SyntheticText2SQL"] |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", required=True, help="package directory, or a Qwen/Qwen3-Embedding-* id for the baseline") |
| ap.add_argument("--tasks", required=True, help="'coir' or a comma-separated list of mteb task names") |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--prompts", help="task prompts JSON (default: <package>/task_prompts.json)") |
| ap.add_argument("--batch", type=int, default=32) |
| ap.add_argument("--efficient-attn", action="store_true", |
| help="repeat K/V instead of SDPA enable_gqa, so fp32 attention uses the memory-efficient kernel " |
| "(needed below ~40 GB of GPU memory; same math, ~1e-7 relative)") |
| ap.add_argument("--token-budget", type=int, default=0, |
| help="replace fixed batches with length-sorted batches of <= N padded tokens (small GPUs)") |
| ap.add_argument("--endpoint", help="encode through a deployed endpoint instead of locally: " |
| "https://<dedicated-dns>/v1/projects/<p>/locations/<r>/endpoints/<id>:predict") |
| a = ap.parse_args() |
| import inspect, mteb, torch |
| from mteb.models.model_implementations.qwen3_models import q3e_instruct_loader |
| if a.efficient_attn: |
| import transformers.integrations.sdpa_attention as sdpa_mod |
| sdpa_mod.use_gqa_in_sdpa = lambda *args, **kwargs: False |
|
|
| out = Path(a.out); out.mkdir(parents=True, exist_ok=True) |
| tasks = COIR if a.tasks.lower() == "coir" else [t.strip() for t in a.tasks.split(",") if t.strip()] |
| local = Path(a.model).is_dir() |
| prompts = json.load(open(a.prompts or Path(a.model) / "task_prompts.json")) |
| if local: |
| dev = {"device": "cpu"} if a.endpoint else {} |
| model = q3e_instruct_loader(a.model, revision=None, prompts_dict=prompts, **dev) |
| prov = json.load(open(Path(a.model) / "PROVENANCE.json")) |
| base_meta = mteb.get_model_meta(prov["base"]["repo"]) |
| model.mteb_model_meta = base_meta.model_copy(update={"name": f"papr-ai/{prov['id']}", "revision": prov["trunk_sha256"][:40]}) |
| else: |
| base_meta = mteb.get_model_meta(a.model) |
| model = q3e_instruct_loader(a.model, revision=base_meta.revision, prompts_dict=prompts, max_seq_length=8192) |
| model.mteb_model_meta = base_meta |
| st = model.model |
| hf = st[0].auto_model |
| if a.token_budget: |
| import numpy as np |
| orig = st.encode |
|
|
| def budget_encode(sentences, prompt=None, batch_size=32, **kw): |
| if isinstance(sentences, str) or len(sentences) == 0: |
| return orig(sentences, prompt=prompt, batch_size=batch_size, **kw) |
| texts = [(prompt or "") + s for s in sentences] |
| lens = [len(x) for x in st.tokenizer(texts, truncation=True, max_length=st.max_seq_length)["input_ids"]] |
| order = sorted(range(len(texts)), key=lambda i: -lens[i]) |
| out, i = [None] * len(texts), 0 |
| while i < len(order): |
| top, j = lens[order[i]], i + 1 |
| while j < len(order) and (j - i + 1) * top <= a.token_budget and j - i < 256: |
| j += 1 |
| idx = order[i:j] |
| embs = orig([sentences[k] for k in idx], prompt=prompt, batch_size=len(idx), **kw) |
| for k, e in zip(idx, embs): |
| out[k] = e |
| i = j |
| return torch.stack(out) if isinstance(out[0], torch.Tensor) else np.stack(out) |
|
|
| st.encode = budget_encode |
| if a.endpoint: |
| import concurrent.futures as cf, threading, urllib.error, urllib.request |
| import google.auth, google.auth.transport.requests |
| import numpy as np |
| creds, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"]) |
| lock = threading.Lock() |
|
|
| def token() -> str: |
| with lock: |
| if not creds.valid: |
| creds.refresh(google.auth.transport.requests.Request()) |
| return creds.token |
|
|
| def call(chunk): |
| body = json.dumps({"instances": [{"inputs": t} for t in chunk]}).encode() |
| for attempt in range(6): |
| req = urllib.request.Request(a.endpoint, data=body, headers={ |
| "Authorization": f"Bearer {token()}", "Content-Type": "application/json"}) |
| try: |
| with urllib.request.urlopen(req, timeout=900) as r: |
| return json.loads(r.read())["predictions"] |
| except urllib.error.HTTPError as e: |
| if e.code in (429, 500, 502, 503, 504) and attempt < 5: |
| time.sleep(2 ** attempt) |
| continue |
| raise |
|
|
| def endpoint_encode(sentences, prompt=None, batch_size=32, **kw): |
| single = isinstance(sentences, str) |
| texts = [(prompt or "") + s for s in ([sentences] if single else sentences)] |
| chunks = [texts[i:i + 32] for i in range(0, len(texts), 32)] |
| with cf.ThreadPoolExecutor(8) as ex: |
| arr = np.asarray([v for part in ex.map(call, chunks) for v in part], dtype=np.float32) |
| return arr[0] if single else arr |
|
|
| st.encode = endpoint_encode |
| info = {"model": a.model, "dtype": str(next(hf.parameters()).dtype), "max_seq_length": st.max_seq_length, |
| "padding_side": st[0].tokenizer.padding_side, "attn": hf.config._attn_implementation, "batch": a.batch, |
| "efficient_attn": a.efficient_attn, "token_budget": a.token_budget, "endpoint": a.endpoint, |
| "tf32": torch.backends.cuda.matmul.allow_tf32, "torch": torch.__version__, "mteb": mteb.__version__} |
| print(f"[load] {info}", flush=True) |
| assert info["dtype"] == "torch.float32" and info["max_seq_length"] == 8192, info |
|
|
| sig = inspect.signature(mteb.evaluate).parameters |
| summary = {} |
| for task in tasks: |
| t0 = time.time() |
| kw = {"encode_kwargs": {"batch_size": a.batch}} |
| if "cache" in sig: |
| RC = getattr(mteb, "ResultCache", None) |
| if RC is None: |
| from mteb.cache import ResultCache as RC |
| kw["cache"] = RC(cache_path=str(out / "_mteb_cache")) |
| if "overwrite_strategy" in sig: |
| kw["overwrite_strategy"] = "always" |
| if "prediction_folder" in sig: |
| kw["prediction_folder"] = str(out / f"_pred_{task}") |
| res = mteb.evaluate(model, [mteb.get_task(task)], **kw) |
| tr = list(getattr(res, "task_results", res))[0] |
| subsets = {f"{split}/{r.get('hf_subset', 'default')}": r.get("ndcg_at_10", r.get("main_score")) |
| for split, rows in (tr.scores or {}).items() for r in rows} |
| vals = [v for v in subsets.values() if v is not None] |
| rec = {"task": task, "ndcg_at_10": sum(vals) / len(vals) if vals else None, "main_score": tr.get_score(), |
| "subsets": subsets, "seconds": round(time.time() - t0), "instruction": prompts.get(task), **info} |
| json.dump(rec, open(out / f"{task}.json", "w"), indent=1) |
| summary[task] = rec["ndcg_at_10"] |
| print(f"[task] {task}: nDCG@10 {rec['ndcg_at_10'] * 100:.2f} ({rec['seconds']}s)", flush=True) |
| json.dump({"info": info, "ndcg_at_10": summary}, open(out / "summary.json", "w"), indent=1) |
| if set(COIR) <= set(summary): |
| print(f"[summary] CoIR-10 mean {sum(summary[t] for t in COIR) / 10 * 100:.2f}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
