lean-RAG

BGE-M3 Mathlib-RAG

README_en:https://huggingface.co/YuxuanGong/lean-RAG/blob/main/README_en.md

基于 BAAI/bge-m3 微调的 Mathlib4 API 检索模型，专为 Lean 4 形式化证明场景设计。该模型能够根据当前证明目标（goal type）和数学关键词，从 Mathlib4 的 32.6 万条声明中检索最相关的引理/定理/定义。

在我的github主页可访问完整skill：https://github.com/Tangtaizong-BUAA/Lean-formulizer

模型用途

该模型是 proof-formalizer 技能链路的检索核心。在形式化一个数学证明时：

1. 从 Lean 环境中提取当前 ⊢ 目标类型（goal conclusion）
2. 结合人类提供的数学关键词
3. 用本模型编码为 1024 维向量
4. 在预构建的 FAISS 索引中检索 top-k 最相关的 Mathlib4 声明

支持中英文混合查询，训练数据包含中文数学描述和 Lean 代码。

使用方式

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("YuxuanGong/lean-RAG")

# 查询：Lean goal + 关键词
query = "x % y < y  nat mod lt"
# 候选：Mathlib4 声明（从 declarations.jsonl 构建的 FAISS 索引）
passages = [
    "Keywords: nat mod lt theorem | Name: Nat.mod_lt | Kind: theorem | ...",
    "Keywords: nat mod add theorem | Name: Nat.ModEq.add | Kind: theorem | ...",
    # ... 32.6 万条
]

query_emb = model.encode(query)
passage_embs = model.encode(passages)
similarities = model.similarity(query_emb, passage_embs)

与 proof-formalizer Skill 集成

# Skill 中的 mathlib_lookup.sh (Route 4) 使用本模型:
DENSE_QUERY="${GOAL_TYPE} ${KEYWORDS_TEXT}"
# 然后调用 search.py 用本模型编码查询 -> FAISS 搜索

模型架构

基于 BGE-M3 (XLM-RoBERTa) 架构：

SentenceTransformer(
  (0): Transformer(XLMRobertaModel, task='feature-extraction')
  (1): Pooling(cls pooling, 1024d)
  (2): Normalize()
)

训练数据

从 Mathlib4 的 325,959 条声明中构造三元组训练数据：

• 训练集: 323,977 条 (query, positive, negatives)
• 验证集: 500 条
• 格式: BGE 三元组格式 — 每条样本包含 1 个 query、1 个 positive passage、7 个 negative passages
• Query: 定理名 + 类型签名（中英混合）
• Positive: 对应引理的完整结构化描述（Keywords | Name | Kind | Namespace | Doc | Type）
• Negatives: 随机采样 + 困难负样本

数据示例

训练细节

训练 Loss 曲线

评测结果

直接检索（Direct Retrieval）

给定声明名 + docstring，检索对应的 Mathlib 声明：

Skill Mode F：引理检索（Lemma Conclusion → Find Lemma）

模拟真实使用场景 — "我有这个 goal，找到能证它的引理"：

Skill Mode E：证明初始检索（Initial Goal → Proof Lemmas）

模拟 "开始证这个定理，需要哪些引理"（更难任务）：

已知限制

• Mode E 召回率较低：从零开始的证明检索仍具有挑战性，MRR 仅 4.7%
• 跨版本 Mathlib 兼容性：模型训练于特定 Mathlib4 版本，API 可能随版本变化
• 最大序列长度：虽然 BGE-M3 支持 8192 tokens，训练时截断至 400 tokens（适配 declaration 长度）
• 中英文偏好：英文查询效果通常优于中文

环境依赖

pip install sentence-transformers>=5.4.1

框架版本（训练时）

• Python: 3.12.3
• Sentence Transformers: 5.4.1
• Transformers: 5.5.4
• PyTorch: 2.8.0+cu128
• Accelerate: 1.13.0

引用

如果你使用了这个模型，请引用：

@misc{gong2025mathlib-rag,
  title={Mathlib-RAG: Fine-tuned BGE-M3 for Lean 4 Mathlib API Retrieval},
  author={Gong, Yuxuan},
  year={2025},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/YuxuanGong/lean-RAG}}
}

基座模型引用：

@article{chen2024bge,
  title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
  author={Chen, Jianlv and Xiao, Shitao and Zhang, Peitian and others},
  journal={arXiv preprint arXiv:2402.03216},
  year={2024}
}

许可

Apache 2.0