Example Scifact Halves

1"""The SKILL.md example: bge-small-en-v1.5 on the first 100 documents of mteb/scifact's corpus, each split into its first
2and second half at the middle sentence boundary, checked by the skill's script with one shared mean. Needs datasets and
3sentence-transformers besides numpy; the inputs are pinned to revisions, and the output records the library versions.
4
5  python example_scifact_halves.py OUTDIR
6
7Sentences are split after '.', '!' or '?' followed by whitespace; the first half is the first len(sentences) // 2 of
8them (at least one), and the second half is the rest.
9"""
10import json
11import os
12import re
13import subprocess
14import sys
15
16import datasets
17import numpy as np
18import sentence_transformers
19import torch
20from datasets import load_dataset
21from sentence_transformers import SentenceTransformer
22
23DATASET, DATASET_REV = "mteb/scifact", "cf10ab6856b15b0e670ef8ae5dae4e266c12d035"
24MODEL, MODEL_REV = "BAAI/bge-small-en-v1.5", "5c38ec7c405ec4b44b94cc5a9bb96e735b38267a"
25
26out = sys.argv[1]
27os.makedirs(out, exist_ok=True)
28docs = load_dataset(DATASET, "corpus", split="corpus", revision=DATASET_REV).select(range(100))
29halves = []
30for text in list(docs["text"]):
31    s = re.split(r"(?<=[.!?])\s+", text.strip())
32    k = max(1, len(s) // 2)
33    halves.append((" ".join(s[:k]), " ".join(s[k:]) or s[-1]))
34m = SentenceTransformer(MODEL, revision=MODEL_REV, device="cpu")
35np.save(os.path.join(out, "A.npy"), m.encode([h[0] for h in halves], normalize_embeddings=True))
36np.save(os.path.join(out, "B.npy"), m.encode([h[1] for h in halves], normalize_embeddings=True))
37here = os.path.dirname(os.path.abspath(__file__))
38script = next(p for p in (os.path.join(here, "anisotropy_check.py"), os.path.join(here, "embedding-anisotropy-check", "scripts", "anisotropy_check.py")) if os.path.exists(p))
39res = subprocess.run([sys.executable, script, os.path.join(out, "A.npy"), os.path.join(out, "B.npy"), "--shared"],
40                     capture_output=True, text=True, check=True).stdout
41res = json.loads(res)
42res["data"] = {"corpus": f"{DATASET}@{DATASET_REV} corpus, rows 0-99", "ids": list(docs["_id"]), "model": f"{MODEL}@{MODEL_REV}",
43               "versions": {"datasets": datasets.__version__, "sentence_transformers": sentence_transformers.__version__,
44                            "torch": torch.__version__, "numpy": np.__version__}}
45json.dump(res, open(os.path.join(out, "example_scifact_halves.json"), "w"), indent=1)
46print(json.dumps({k: res[k] for k in ("anisotropy", "raw", "centred")}))