mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
108 lines
4.6 KiB
Python
108 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Mathlib4-to-Parquet Global Ingester
|
|
Transforms the entire Mathlib4 library into a columnar dataset for large-scale
|
|
Lean 4 training (LeanGPT / SNN).
|
|
|
|
License: Apache 2.0 (Lean Community)
|
|
"""
|
|
|
|
import os
|
|
import hashlib
|
|
from pathlib import Path
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
from datetime import datetime
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# §0 Configuration
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
|
|
RESEARCH_STACK_ROOT = Path("/home/allaun/Research Stack")
|
|
MATHLIB_SRC_ROOT = RESEARCH_STACK_ROOT / "0-Core-Formalism/lean/Semantics/.lake/packages/mathlib/Mathlib"
|
|
OUTPUT_DIR = RESEARCH_STACK_ROOT / "shared-data/data/datasets"
|
|
OUTPUT_FILE = OUTPUT_DIR / "mathlib4_complete.parquet"
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# §1 Extraction Logic
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
|
|
def get_mathlib_domain(file_path: Path) -> str:
|
|
"""Extracts the top-level math domain (e.g., 'Analysis')."""
|
|
try:
|
|
relative = file_path.relative_to(MATHLIB_SRC_ROOT)
|
|
return relative.parts[0] if len(relative.parts) > 1 else "Core"
|
|
except:
|
|
return "Unknown"
|
|
|
|
def ingest_mathlib():
|
|
"""Gathers Mathlib files with progress tracking."""
|
|
rows = []
|
|
|
|
if not MATHLIB_SRC_ROOT.exists():
|
|
print(f"[ERROR] Mathlib root not found: {MATHLIB_SRC_ROOT}")
|
|
return []
|
|
|
|
print(f"[INFO] Scanning Mathlib corpus at {MATHLIB_SRC_ROOT}...")
|
|
lean_files = list(MATHLIB_SRC_ROOT.rglob("*.lean"))
|
|
total_files = len(lean_files)
|
|
print(f"[INFO] Found {total_files} Lean files.")
|
|
|
|
for i, lean_file in enumerate(lean_files):
|
|
try:
|
|
with open(lean_file, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
|
|
# Module name (relative path without .lean)
|
|
module_name = str(lean_file.relative_to(MATHLIB_SRC_ROOT)).replace("/", ".").replace(".lean", "")
|
|
domain = get_mathlib_domain(lean_file)
|
|
|
|
rows.append({
|
|
"module_name": module_name,
|
|
"domain": domain,
|
|
"content": content,
|
|
"token_count": len(content.split()),
|
|
"hash": hashlib.sha256(content.encode()).hexdigest(),
|
|
"extracted_at": datetime.utcnow().isoformat()
|
|
})
|
|
|
|
if (i + 1) % 500 == 0:
|
|
print(f"[PROGRESS] Ingested {i+1}/{total_files} files...")
|
|
|
|
except Exception as e:
|
|
print(f"[WARN] Failed to process {lean_file}: {e}")
|
|
|
|
return rows
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# §2 Parquet Serialization
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
|
|
def write_parquet(rows):
|
|
"""Writes to Parquet with Schema enforcement."""
|
|
if not rows:
|
|
return
|
|
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
fields = [
|
|
pa.field('module_name', pa.string()),
|
|
pa.field('domain', pa.string()),
|
|
pa.field('content', pa.string()),
|
|
pa.field('token_count', pa.int32()),
|
|
pa.field('hash', pa.string()),
|
|
pa.field('extracted_at', pa.string())
|
|
]
|
|
schema = pa.schema(fields)
|
|
|
|
# Building table
|
|
table = pa.Table.from_pylist(rows, schema=schema)
|
|
|
|
# Writing with compression for such a large file
|
|
pq.write_table(table, OUTPUT_FILE, compression='snappy')
|
|
print(f"\n[SUCCESS] Global Mathlib4 Dataset created: {OUTPUT_FILE}")
|
|
print(f"[INFO] Record count: {len(rows)}")
|
|
|
|
if __name__ == "__main__":
|
|
data = ingest_mathlib()
|
|
write_parquet(data)
|