File size: 6,666 Bytes
032a961 | 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 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | #!/usr/bin/env python3
"""Build the SPIRA training corpus: mixed natural language + source code.
Sources:
- English books from Project Gutenberg (public domain)
- French books from Project Gutenberg (flavour / multilingual robustness)
- Python standard library source files (code domain)
Outputs (under python/data/ next to the scripts):
- train.bin : ~95% of the interleaved stream (text + code)
- test.bin : held-out tail slice (~5%)
- oneshot.bin : a full book never used in training (one-shot / fresh-text tests)
- corpus_manifest.json
"""
import glob
import json
import os
import shutil
import subprocess
import sys
BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = f"{BASE}/data"
BOOKS = f"{DATA}/corpus/books"
CODE = f"{DATA}/corpus/code"
# (gutenberg_id, name, holdout?)
ENGLISH_BOOKS = [
(11, "alice", False),
(84, "frankenstein", False),
(120, "treasure_island", False),
(174, "dorian_gray", False),
(98, "two_cities", False),
(345, "dracula", False),
(5200, "metamorphosis", True), # held out entirely -> oneshot.bin
(1342, "pride_prejudice", False),
(1400, "great_expectations", False),
(1661, "sherlock", False),
(2701, "moby_dick", False),
]
FRENCH_BOOKS = [
(19942, "candide_fr", False),
(17489, "les_mis_fr", False),
]
CODE_BUDGET_BYTES = 6 * 1024 * 1024
SKIP_CODE_DIRS = {"test", "tests", "__pycache__", "lib2to3", "idlelib", "site-packages", "turtledemo"}
SKIP_CODE_FILES = {"setup.py"}
DOC_SEP = b"\n\n"
def strip_gutenberg_header(raw: bytes) -> bytes:
"""Remove the Project Gutenberg boilerplate header/footer when present."""
marker = b"*** START OF"
idx = raw.find(marker)
if idx != -1:
end = raw.find(b"\n", idx)
if end != -1:
raw = raw[end + 1:]
marker = b"*** END OF"
idx = raw.find(marker)
if idx != -1:
raw = raw[:idx]
return raw
def french_ratio(raw: bytes) -> float:
"""Fraction of letters that carry French diacritics (cheap language check)."""
letters = 0
for b in raw:
if 65 <= b <= 90 or 97 <= b <= 122:
letters += 1
accented = raw.count(0xC3) # UTF-8 lead byte for most accented Latin letters
if letters == 0:
return 0.0
return accented / max(letters, 1)
def fetch_book(gid: int, name: str) -> str | None:
os.makedirs(BOOKS, exist_ok=True)
dst = f"{BOOKS}/{name}.txt"
if os.path.exists(dst) and os.path.getsize(dst) > 20_000:
return dst
urls = [
f"https://www.gutenberg.org/cache/epub/{gid}/pg{gid}.txt",
f"https://www.gutenberg.org/files/{gid}/{gid}-0.txt",
f"https://www.gutenberg.org/files/{gid}/{gid}.txt",
]
for url in urls:
r = subprocess.run(["curl", "-sL", "-m", "120", "-o", dst, url])
if r.returncode == 0 and os.path.exists(dst) and os.path.getsize(dst) > 20_000:
return dst
if os.path.exists(dst):
os.remove(dst)
return None
def collect_code_files(stdlib_dir: str, budget: int) -> list[str]:
files = []
total = 0
for root, dirs, fnames in os.walk(stdlib_dir):
dirs[:] = [d for d in sorted(dirs) if d not in SKIP_CODE_DIRS]
for fn in sorted(fnames):
if not fn.endswith(".py") or fn in SKIP_CODE_FILES:
continue
path = os.path.join(root, fn)
try:
sz = os.path.getsize(path)
except OSError:
continue
if sz < 200: # skip stubs
continue
files.append(path)
total += sz
if total >= budget:
return files
return files
def main() -> None:
os.makedirs(DATA, exist_ok=True)
os.makedirs(CODE, exist_ok=True)
docs = [] # list of (domain, bytes)
holdout_docs = []
print("== Downloading books ==")
for gid, name, holdout in ENGLISH_BOOKS + FRENCH_BOOKS:
path = fetch_book(gid, name)
if path is None:
print(f" !! failed: {name} (gid={gid})")
continue
raw = open(path, "rb").read()
raw = strip_gutenberg_header(raw)
if len(raw) < 20_000:
print(f" !! too small after strip: {name}")
continue
domain = "fr" if name.endswith("_fr") else "en"
if holdout:
holdout_docs.append((domain, raw))
else:
docs.append((domain, raw))
print(f" ok: {name:18s} {len(raw)/1024:8.1f} KB ({'holdout' if holdout else 'train'})")
print("== Collecting Python stdlib code ==")
import sysconfig
stdlib_dir = sysconfig.get_paths()["stdlib"]
code_files = collect_code_files(stdlib_dir, CODE_BUDGET_BYTES)
print(f" {len(code_files)} files from {stdlib_dir}")
for path in code_files:
try:
raw = open(path, "rb").read()
except OSError:
continue
if not raw:
continue
docs.append(("code", raw))
dst = os.path.join(CODE, os.path.basename(path))
if not os.path.exists(dst):
shutil.copyfile(path, dst)
# ---- interleave documents deterministically (seeded shuffle) ----
import random
rng = random.Random(20261022)
rng.shuffle(docs)
stream = DOC_SEP.join(d[1] for d in docs)
n_test = int(len(stream) * 0.05)
# cut at a document boundary close to the tail
cut = len(stream) - n_test
boundary = stream.rfind(DOC_SEP, 0, cut + n_test // 2)
if boundary == -1 or boundary < len(stream) // 2:
boundary = cut
train = stream[:boundary]
test = stream[boundary + len(DOC_SEP):]
with open(f"{DATA}/train.bin", "wb") as f:
f.write(train)
with open(f"{DATA}/test.bin", "wb") as f:
f.write(test)
oneshot = DOC_SEP.join(d[1] for d in holdout_docs)
with open(f"{DATA}/oneshot.bin", "wb") as f:
f.write(oneshot)
en_bytes = sum(len(d[1]) for d in docs if d[0] == "en")
fr_bytes = sum(len(d[1]) for d in docs if d[0] == "fr")
code_bytes = sum(len(d[1]) for d in docs if d[0] == "code")
manifest = {
"train_bytes": len(train),
"test_bytes": len(test),
"oneshot_bytes": len(oneshot),
"composition": {"english": en_bytes, "french": fr_bytes, "code": code_bytes},
"n_documents": len(docs),
"n_holdout_documents": len(holdout_docs),
"code_files": len(code_files),
}
with open(f"{DATA}/corpus_manifest.json", "w") as f:
json.dump(manifest, f, indent=2)
print(json.dumps(manifest, indent=2))
print("DONE")
if __name__ == "__main__":
main()
|