embeddinggemma-2-lean / scripts /lean_lemma_retriever.py

Upload lean_lemma_retriever.py

ff241f2 verified 25 days ago

8.16 kB

#!/usr/bin/env python3
"""
lean_lemma_retriever.py — EmbeddingGemma-2 Lean 引理检索器
接入 scientific-discovery-proof 技能 Stage 3 引理召回 (v2.22.0+)。
用途: 给定 Lean 定理名(或自由文本), 召回其依赖 DAG 中可能引用的候选引理/定义。
基座: google/embeddinggemma-2-740m, 在 openai_math Lean 依赖对(377K)上全参微调。
recall@10 = 0.5505 (增强文本口径: 名字词化 + 类型签名120 + docstring60)。
运行依赖(独立 venv, 勿用 Hermes 主 venv):
    /mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/venv/bin/python
    torch 2.5.1+cu121 / transformers 5.19.0 / sentence-transformers 6.1.0
用法:
    # 建索引(首次, 编码 18 万声明)
    python lean_lemma_retriever.py --build-index
    # 召回
    python lean_lemma_retriever.py --query resolutionCokernelMap --top-k 10
    # 自由文本召回
    python lean_lemma_retriever.py --text "spectral gap of Laplacian" --top-k 10
"""
import re, json, os, sys, argparse

MODEL_DIR = "/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/embeddinggemma-2-lean"
DECLS_PATH = "/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/declarations.jsonl"
INDEX_DIR = "/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/index"


# ---------------------------------------------------------------- 文本口径

def name_to_text(name: str) -> str:
    """Lean 标识符 → 可读文本 (与训练口径一致)"""
    s = re.sub(r'_+', ' ', name)
    s = re.sub(r'([a-z0-9])([A-Z])', r'\1 \2', s)
    return s.strip()


def enrich(name: str, decls: dict, sig_chars: int = 120, doc_chars: int = 60) -> str:
    """增强文本: 名字词化 | 类型签名 | docstring (与训练口径严格一致)"""
    parts = [name_to_text(name)]
    d = decls.get(name)
    if d:
        if d.get("signature"):
            parts.append("type " + d["signature"][:sig_chars])
        if d.get("docstring"):
            parts.append(d["docstring"][:doc_chars])
    return " | ".join(parts)


def load_decls(path: str = DECLS_PATH) -> dict:
    decls = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            d = json.loads(line)
            decls[d["name"]] = d
    return decls


# ---------------------------------------------------------------- 模型与检索

class LemmaRetriever:
    """把 Lean 声明库编码为索引, 供 Stage 3 召回候选引理"""

    def __init__(self, model, names, decls, normalize: bool = True):
        self.model = model
        self.names = list(names)
        self.decls = decls
        self.normalize = normalize
        self.emb = None

    def build_index(self, batch_size: int = 64, cache_dir: str = INDEX_DIR):
        import numpy as np
        txts = [enrich(n, self.decls) for n in self.names]
        # 分批编码到 CPU（P620 4GB 显存下 18 万声明全量编码会 OOM，须逐批累积）
        n = len(txts)
        emb_batches = []
        for i in range(0, n, batch_size):
            chunk = txts[i:i + batch_size]
            e = self.model.encode(
                chunk, normalize_embeddings=self.normalize,
                batch_size=batch_size, convert_to_numpy=True, show_progress_bar=False,
            )
            emb_batches.append(e)
            if (i // batch_size) % 500 == 0 and i > 0:
                print(f"    [index] 已编码 {i}/{n}", flush=True)
        emb = np.vstack(emb_batches)
        self.emb = __import__("torch").from_numpy(emb)
        # 缓存
        os.makedirs(cache_dir, exist_ok=True)
        np.save(os.path.join(cache_dir, "emb.npy"), emb)
        with open(os.path.join(cache_dir, "names.json"), "w", encoding="utf-8") as f:
            json.dump(self.names, f)
        return self.emb

    def load_index(self, cache_dir: str = INDEX_DIR):
        import numpy as np
        import torch
        emb = np.load(os.path.join(cache_dir, "emb.npy"))
        with open(os.path.join(cache_dir, "names.json"), encoding="utf-8") as f:
            self.names = json.load(f)
        self.emb = torch.from_numpy(emb)
        return self.emb

    def _encode_query(self, text: str):
        return self.model.encode(
            [text], normalize_embeddings=self.normalize, convert_to_tensor=True
        )

    def recall(self, theorem_name: str, top_k: int = 10):
        """给定定理名, 召回候选依赖 (用增强文本口径)"""
        q = self._encode_query(enrich(theorem_name, self.decls))
        return self._topk(q, top_k)

    def recall_by_text(self, query_text: str, top_k: int = 10):
        """自由文本召回"""
        q = self._encode_query(query_text)
        return self._topk(q, top_k)

    def _topk(self, q, top_k: int):
        import torch
        # 统一到 CPU 做矩阵乘法（18万×768 的 matmul CPU 上仅数百毫秒）
        q_cpu = q.cpu() if hasattr(q, "is_cuda") and q.is_cuda else q
        emb = self.emb.cpu() if hasattr(self.emb, "is_cuda") and self.emb.is_cuda else self.emb
        sim = (q_cpu @ emb.T)[0]
        k = min(top_k, len(self.names))
        idx = torch.topk(sim, k=k).indices.tolist()
        return [(self.names[i], float(sim[i])) for i in idx]


def load_model(device: str = "cpu", dtype: str = "float32"):
    """加载模型. P620(Pascal) 不支持 bf16, 默认 float32."""
    from sentence_transformers import SentenceTransformer
    kwargs = {}
    if dtype == "float32":
        import torch
        kwargs["torch_dtype"] = torch.float32
    model = SentenceTransformer(MODEL_DIR, device=device, model_kwargs=kwargs)
    model.max_seq_length = 64
    return model


# ---------------------------------------------------------------- CLI

def main():
    global MODEL_DIR, DECLS_PATH
    ap = argparse.ArgumentParser(description="EmbeddingGemma-2 Lean 引理检索器")
    ap.add_argument("--model-dir", default=MODEL_DIR)
    ap.add_argument("--decls", default=DECLS_PATH)
    ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
    ap.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"])
    ap.add_argument("--build-index", action="store_true", help="编码声明库并缓存")
    ap.add_argument("--query", type=str, help="定理名召回")
    ap.add_argument("--text", type=str, help="自由文本召回")
    ap.add_argument("--top-k", type=int, default=10)
    args = ap.parse_args()

    MODEL_DIR = args.model_dir
    DECLS_PATH = args.decls

    print(f"[load] 加载模型 {MODEL_DIR} (device={args.device}, dtype={args.dtype})", flush=True)
    model = load_model(args.device, args.dtype)
    print("[load] 模型加载完成", flush=True)

    print(f"[load] 加载声明库 {DECLS_PATH}", flush=True)
    decls = load_decls(args.decls)
    names = list(decls.keys())
    print(f"[load] 声明库 {len(names)} 条", flush=True)

    retriever = LemmaRetriever(model, names, decls)

    if args.build_index:
        print("[index] 编码声明库...", flush=True)
        retriever.build_index()
        print(f"[index] 完成, 缓存到 {INDEX_DIR}", flush=True)
        return

    # 尝试加载缓存索引
    import os as _os
    emb_path = _os.path.join(INDEX_DIR, "emb.npy")
    if _os.path.exists(emb_path):
        print("[index] 加载缓存索引...", flush=True)
        retriever.load_index()
    else:
        print("[index] 无缓存, 编码声明库 (约 2 分钟 GPU / 30+ 分钟 CPU)...", flush=True)
        retriever.build_index()

    if args.query:
        print(f"[recall] 定理名: {args.query}", flush=True)
        for name, score in retriever.recall(args.query, args.top_k):
            d = decls.get(name, {})
            kind = d.get("kind", "?")
            print(f"  {score:.4f}  [{kind}] {name}")
    elif args.text:
        print(f"[recall] 文本: {args.text}", flush=True)
        for name, score in retriever.recall_by_text(args.text, args.top_k):
            d = decls.get(name, {})
            kind = d.get("kind", "?")
            print(f"  {score:.4f}  [{kind}] {name}")
    else:
        print("用法: 指定 --query 或 --text")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	lean_lemma_retriever.py — EmbeddingGemma-2 Lean 引理检索器

	接入 scientific-discovery-proof 技能 Stage 3 引理召回 (v2.22.0+)。

	用途: 给定 Lean 定理名(或自由文本), 召回其依赖 DAG 中可能引用的候选引理/定义。
	基座: google/embeddinggemma-2-740m, 在 openai_math Lean 依赖对(377K)上全参微调。
	recall@10 = 0.5505 (增强文本口径: 名字词化 + 类型签名120 + docstring60)。

	运行依赖(独立 venv, 勿用 Hermes 主 venv):
	/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/venv/bin/python
	torch 2.5.1+cu121 / transformers 5.19.0 / sentence-transformers 6.1.0

	用法:
	# 建索引(首次, 编码 18 万声明)
	python lean_lemma_retriever.py --build-index
	# 召回
	python lean_lemma_retriever.py --query resolutionCokernelMap --top-k 10
	# 自由文本召回
	python lean_lemma_retriever.py --text "spectral gap of Laplacian" --top-k 10
	"""
	import re, json, os, sys, argparse

	MODEL_DIR = "/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/embeddinggemma-2-lean"
	DECLS_PATH = "/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/declarations.jsonl"
	INDEX_DIR = "/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/index"


	# ---------------------------------------------------------------- 文本口径

	def name_to_text(name: str) -> str:
	"""Lean 标识符 → 可读文本 (与训练口径一致)"""
	s = re.sub(r'_+', ' ', name)
	s = re.sub(r'([a-z0-9])([A-Z])', r'\1 \2', s)
	return s.strip()


	def enrich(name: str, decls: dict, sig_chars: int = 120, doc_chars: int = 60) -> str:
	"""增强文本: 名字词化 \| 类型签名 \| docstring (与训练口径严格一致)"""
	parts = [name_to_text(name)]
	d = decls.get(name)
	if d:
	if d.get("signature"):
	parts.append("type " + d["signature"][:sig_chars])
	if d.get("docstring"):
	parts.append(d["docstring"][:doc_chars])
	return " \| ".join(parts)


	def load_decls(path: str = DECLS_PATH) -> dict:
	decls = {}
	with open(path, encoding="utf-8") as f:
	for line in f:
	line = line.strip()
	if not line:
	continue
	d = json.loads(line)
	decls[d["name"]] = d
	return decls


	# ---------------------------------------------------------------- 模型与检索

	class LemmaRetriever:
	"""把 Lean 声明库编码为索引, 供 Stage 3 召回候选引理"""

	def __init__(self, model, names, decls, normalize: bool = True):
	self.model = model
	self.names = list(names)
	self.decls = decls
	self.normalize = normalize
	self.emb = None

	def build_index(self, batch_size: int = 64, cache_dir: str = INDEX_DIR):
	import numpy as np
	txts = [enrich(n, self.decls) for n in self.names]
	# 分批编码到 CPU（P620 4GB 显存下 18 万声明全量编码会 OOM，须逐批累积）
	n = len(txts)
	emb_batches = []
	for i in range(0, n, batch_size):
	chunk = txts[i:i + batch_size]
	e = self.model.encode(
	chunk, normalize_embeddings=self.normalize,
	batch_size=batch_size, convert_to_numpy=True, show_progress_bar=False,
	)
	emb_batches.append(e)
	if (i // batch_size) % 500 == 0 and i > 0:
	print(f" [index] 已编码 {i}/{n}", flush=True)
	emb = np.vstack(emb_batches)
	self.emb = __import__("torch").from_numpy(emb)
	# 缓存
	os.makedirs(cache_dir, exist_ok=True)
	np.save(os.path.join(cache_dir, "emb.npy"), emb)
	with open(os.path.join(cache_dir, "names.json"), "w", encoding="utf-8") as f:
	json.dump(self.names, f)
	return self.emb

	def load_index(self, cache_dir: str = INDEX_DIR):
	import numpy as np
	import torch
	emb = np.load(os.path.join(cache_dir, "emb.npy"))
	with open(os.path.join(cache_dir, "names.json"), encoding="utf-8") as f:
	self.names = json.load(f)
	self.emb = torch.from_numpy(emb)
	return self.emb

	def _encode_query(self, text: str):
	return self.model.encode(
	[text], normalize_embeddings=self.normalize, convert_to_tensor=True
	)

	def recall(self, theorem_name: str, top_k: int = 10):
	"""给定定理名, 召回候选依赖 (用增强文本口径)"""
	q = self._encode_query(enrich(theorem_name, self.decls))
	return self._topk(q, top_k)

	def recall_by_text(self, query_text: str, top_k: int = 10):
	"""自由文本召回"""
	q = self._encode_query(query_text)
	return self._topk(q, top_k)

	def _topk(self, q, top_k: int):
	import torch
	# 统一到 CPU 做矩阵乘法（18万×768 的 matmul CPU 上仅数百毫秒）
	q_cpu = q.cpu() if hasattr(q, "is_cuda") and q.is_cuda else q
	emb = self.emb.cpu() if hasattr(self.emb, "is_cuda") and self.emb.is_cuda else self.emb
	sim = (q_cpu @ emb.T)[0]
	k = min(top_k, len(self.names))
	idx = torch.topk(sim, k=k).indices.tolist()
	return [(self.names[i], float(sim[i])) for i in idx]


	def load_model(device: str = "cpu", dtype: str = "float32"):
	"""加载模型. P620(Pascal) 不支持 bf16, 默认 float32."""
	from sentence_transformers import SentenceTransformer
	kwargs = {}
	if dtype == "float32":
	import torch
	kwargs["torch_dtype"] = torch.float32
	model = SentenceTransformer(MODEL_DIR, device=device, model_kwargs=kwargs)
	model.max_seq_length = 64
	return model


	# ---------------------------------------------------------------- CLI

	def main():
	global MODEL_DIR, DECLS_PATH
	ap = argparse.ArgumentParser(description="EmbeddingGemma-2 Lean 引理检索器")
	ap.add_argument("--model-dir", default=MODEL_DIR)
	ap.add_argument("--decls", default=DECLS_PATH)
	ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
	ap.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"])
	ap.add_argument("--build-index", action="store_true", help="编码声明库并缓存")
	ap.add_argument("--query", type=str, help="定理名召回")
	ap.add_argument("--text", type=str, help="自由文本召回")
	ap.add_argument("--top-k", type=int, default=10)
	args = ap.parse_args()

	MODEL_DIR = args.model_dir
	DECLS_PATH = args.decls

	print(f"[load] 加载模型 {MODEL_DIR} (device={args.device}, dtype={args.dtype})", flush=True)
	model = load_model(args.device, args.dtype)
	print("[load] 模型加载完成", flush=True)

	print(f"[load] 加载声明库 {DECLS_PATH}", flush=True)
	decls = load_decls(args.decls)
	names = list(decls.keys())
	print(f"[load] 声明库 {len(names)} 条", flush=True)

	retriever = LemmaRetriever(model, names, decls)

	if args.build_index:
	print("[index] 编码声明库...", flush=True)
	retriever.build_index()
	print(f"[index] 完成, 缓存到 {INDEX_DIR}", flush=True)
	return

	# 尝试加载缓存索引
	import os as _os
	emb_path = _os.path.join(INDEX_DIR, "emb.npy")
	if _os.path.exists(emb_path):
	print("[index] 加载缓存索引...", flush=True)
	retriever.load_index()
	else:
	print("[index] 无缓存, 编码声明库 (约 2 分钟 GPU / 30+ 分钟 CPU)...", flush=True)
	retriever.build_index()

	if args.query:
	print(f"[recall] 定理名: {args.query}", flush=True)
	for name, score in retriever.recall(args.query, args.top_k):
	d = decls.get(name, {})
	kind = d.get("kind", "?")
	print(f" {score:.4f} [{kind}] {name}")
	elif args.text:
	print(f"[recall] 文本: {args.text}", flush=True)
	for name, score in retriever.recall_by_text(args.text, args.top_k):
	d = decls.get(name, {})
	kind = d.get("kind", "?")
	print(f" {score:.4f} [{kind}] {name}")
	else:
	print("用法: 指定 --query 或 --text")


	if __name__ == "__main__":
	main()