File size: 5,787 Bytes
d7c04c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | #!/usr/bin/env python3
"""
parse_lean_declarations.py — 从 Lean 源码解析声明库 declarations.jsonl
输入: Lean 源码根目录 (openai_math-main/lean, ~150K .lean 文件)
输出: declarations.jsonl — 每行 {"name": 短名, "kind": ..., "signature": ..., "docstring": ...}
与 README_接入方法.md 第五节口径一致:
- 正则匹配 (theorem|lemma|def|abbrev|instance|structure|class|inductive|opaque) 名字
- 签名取到第一个 := 或 where 前
- docstring 取声明上方紧邻的 /-- ... -/ 块
- 短名 = 点分名最后一段; 冲突保留首见
"""
import re, json, os, sys, time
LEAN_ROOT = sys.argv[1] if len(sys.argv) > 1 else "/mnt/d/AI_for_Science/openai_math-main/lean"
OUT = sys.argv[2] if len(sys.argv) > 2 else "/mnt/d/AI_Training/EmbeddingGemma2_Lean检索模型_v2/declarations.jsonl"
KINDS = ("theorem", "lemma", "def", "abbrev", "instance", "structure", "class", "inductive", "opaque")
# 声明起始: 行首(允许空白) kind 空格 标识符
DECL_RE = re.compile(
r'^\s*(theorem|lemma|def|abbrev|instance|structure|class|inductive|opaque)\s+'
r'([A-Za-z_α-ωΑ-Ω0-9][\wα-ωΑ-Ω.\'’]*)\b'
)
# docstring 块: /-- ... -/ (多行)
DOC_OPEN = re.compile(r'^\s*/--')
DOC_CLOSE = re.compile(r'-/\s*$')
def short_name(full: str) -> str:
return full.split('.')[-1]
def collect_files(root):
for dirpath, _, filenames in os.walk(root):
for fn in filenames:
if fn.endswith('.lean'):
yield os.path.join(dirpath, fn)
def parse_file(path):
"""返回该文件内 [(name, kind, signature, docstring), ...]"""
out = []
try:
with open(path, encoding='utf-8', errors='replace') as f:
lines = f.readlines()
except Exception:
return out
n = len(lines)
# 预扫描: 每个 docstring 块的结束行 -> 内容
# 用简单状态机记录最近的 docstring 结束行
doc_end_line = -1 # 最近一个 docstring 块的结束行索引
doc_text = ""
in_doc = False
doc_buf = []
doc_end_line = -1
i = 0
# 先收集 docstring 结束位置
doc_ends = {} # line_idx -> docstring text (不含 /-- 和 -/)
j = 0
while j < n:
line = lines[j]
if not in_doc and DOC_OPEN.match(line):
in_doc = True
doc_buf = []
# 同行闭合?
rest = line
# 去掉开头的 /--
m = DOC_OPEN.match(line)
# 收集这行 /-- 之后的内容
content = line[m.end():]
if DOC_CLOSE.search(content):
# 单行 docstring
content = DOC_CLOSE.sub('', content).strip()
doc_buf.append(content)
doc_ends[j] = ' '.join(doc_buf).strip()
in_doc = False
doc_buf = []
else:
doc_buf.append(content.strip())
elif in_doc:
if DOC_CLOSE.search(line):
content = DOC_CLOSE.sub('', line).strip()
if content:
doc_buf.append(content)
doc_ends[j] = ' '.join(x for x in doc_buf if x).strip()
in_doc = False
doc_buf = []
else:
doc_buf.append(line.strip())
j += 1
# 扫描声明
k = 0
while k < n:
m = DECL_RE.match(lines[k])
if not m:
k += 1
continue
kind = m.group(1)
full_name = m.group(2)
name = short_name(full_name)
# docstring: 找声明行上方最近的 docstring 结束行 (空行容忍)
doc = ""
for back in range(k-1, max(k-4, -1), -1):
if back in doc_ends:
doc = doc_ends[back]
break
if lines[back].strip() != "":
# 遇到非空非 docstring 行，停止
break
# 签名: 从声明名之后累积到第一个 := 或 where 或换行
sig_parts = []
# 当前行名字之后的内容
rest = lines[k][m.end():]
sig_parts.append(rest.rstrip('\n'))
# 若当前行没有 := / where，继续向下累积
joined = rest
kk = k + 1
if not re.search(r':=|\bwhere\b', joined):
while kk < n:
joined += lines[kk]
sig_parts.append(lines[kk].rstrip('\n'))
if re.search(r':=|\bwhere\b', lines[kk]):
break
kk += 1
if kk - k > 40: # 保护: 签名过长
break
signature = ' '.join(sig_parts)
signature = re.split(r':=|\bwhere\b', signature)[0].strip()
# 压缩空白
signature = re.sub(r'\s+', ' ', signature).strip()
out.append((name, kind, signature, doc))
k += 1
return out
def main():
t0 = time.time()
seen = {} # 短名 -> 首见声明
files = list(collect_files(LEAN_ROOT))
total = len(files)
print(f"[parse] 共 {total} 个 .lean 文件", flush=True)
count = 0
for idx, path in enumerate(files):
for (name, kind, sig, doc) in parse_file(path):
if name not in seen:
seen[name] = {"name": name, "kind": kind, "signature": sig, "docstring": doc}
count += 1
if count % 20000 == 0:
print(f"[parse] 进度 {count}/{total} ({time.time()-t0:.0f}s), 声明累计 {len(seen)}", flush=True)
# 写出
os.makedirs(os.path.dirname(OUT), exist_ok=True)
with open(OUT, 'w', encoding='utf-8') as f:
for d in seen.values():
f.write(json.dumps(d, ensure_ascii=False) + '\n')
print(f"[parse] DONE: {len(seen)} 声明 -> {OUT} ({time.time()-t0:.0f}s)", flush=True)
if __name__ == '__main__':
main()
|