File size: 2,143 Bytes
4397e12 | 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 | """Shared text constants and raw-document readers."""
from __future__ import annotations
import glob
import gzip
import json
import os
import pyarrow.parquet as pq
DATA = os.environ.get("TA_DATA", "data")
# Special tokens get ids 0..len-1 in this order. Hermes-style tool tags so the later Hermes
# phase only has to change the system prompt and tool schema, not the token vocabulary.
SPECIAL_TOKENS = [
"<|endoftext|>", "<|pad|>", "<|im_start|>", "<|im_end|>",
"<think>", "</think>", "<tool_call>", "</tool_call>", "<tool_response>", "</tool_response>",
] + [f"<|reserved_{i}|>" for i in range(22)]
EOS_ID = 0
PAD_ID = 1
# stack-edu shard -> language (see scripts/download.py)
CODE_LANG = {}
for lang, shards in {"markdown": [0, 1], "python": [73, 74, 75, 76], "shell": [87, 88], "javascript": [50],
"typescript": [91], "go": [34], "rust": [86], "c": [18]}.items():
for i in shards:
CODE_LANG[i] = lang
def code_files(lang: str) -> list[str]:
return [f"{DATA}/raw/code/stack-edu-{i:04d}.json.gz" for i, l in sorted(CODE_LANG.items()) if l == lang
and os.path.exists(f"{DATA}/raw/code/stack-edu-{i:04d}.json.gz")]
def iter_code(path: str):
with gzip.open(path, "rt") as f:
for line in f:
r = json.loads(line)
p = (r.get("metadata") or {}).get("path") or ""
yield (f"File: {p.lstrip('/')}\n\n" if p else "") + r["text"]
def iter_parquet(path: str):
pf = pq.ParquetFile(path)
for rg in range(pf.num_row_groups):
for t in pf.read_row_group(rg, columns=["text"]).column("text").to_pylist():
yield t
def source_files() -> dict[str, list[tuple[str, callable]]]:
"""Source name -> list of (file, iterator fn)."""
out = {
"english": [(f, iter_parquet) for f in sorted(glob.glob(f"{DATA}/raw/english/**/*.parquet", recursive=True))],
"math": [(f, iter_parquet) for f in sorted(glob.glob(f"{DATA}/raw/math/**/*.parquet", recursive=True))],
}
for lang in sorted(set(CODE_LANG.values())):
out[f"code_{lang}"] = [(f, iter_code) for f in code_files(lang)]
return out
|