""" Chunk quality-filtered documents into: 1. A pre-training corpus (data/pretrain/corpus_pretrain.txt) - Concatenated source code with document separators - Used for base LM pre-training 2. Retrieval chunks (data/chunks/chunks.jsonl) - Semantically meaningful code units (functions, classes, structs) - Used by the search agent and for SFT trace generation Input: data/quality/documents_quality.jsonl Output: data/pretrain/corpus_pretrain.txt data/chunks/chunks.jsonl """ import json import os import re PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) DATA_DIR = os.path.join(PROJECT_DIR, "data") QUALITY_DOCS_PATH = os.path.join(DATA_DIR, "quality", "documents_quality.jsonl") PRETRAIN_DIR = os.path.join(DATA_DIR, "pretrain") CORPUS_PRETRAIN_PATH = os.path.join(PRETRAIN_DIR, "corpus_pretrain.txt") CHUNKS_DIR = os.path.join(DATA_DIR, "chunks") CHUNKS_PATH = os.path.join(CHUNKS_DIR, "chunks.jsonl") # ─── Chunk start patterns per language ─────────────────────────────────────── CHUNK_STARTS = { "python": re.compile(r"^(def |class |@|async def )", re.M), "javascript": re.compile(r"^(function |const |let |class |export |async function|interface |type \w+ =)", re.M), "typescript": re.compile(r"^(function |const |let |class |export |async function|interface |type \w+ =|enum )", re.M), "rust": re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |macro_rules!|pub struct|pub enum|pub trait)", re.M), "go": re.compile(r"^(func |type \w+ struct|type \w+ interface)", re.M), "c": re.compile(r"^(static |void |int |char |size_t |typedef |struct \w+ \{|#define|#if|#ifndef|ngx_)", re.M), "cpp": re.compile(r"^(static |void |int |char |size_t |template |class \w+|namespace |struct \w+|#define|#if|#ifndef)", re.M), "java": re.compile(r"^(public |private |protected |class \w+|interface |enum )", re.M), "csharp": re.compile(r"^(public |private |protected |internal |class \w+|interface |enum )", re.M), "ruby": re.compile(r"^(def |class |module )", re.M), } NAME_PATTERNS = [ (re.compile(r"def (\w+)"), "function"), (re.compile(r"class (\w+)"), "class"), (re.compile(r"fn (\w+)"), "function"), (re.compile(r"pub fn (\w+)"), "function"), (re.compile(r"struct (\w+)"), "struct"), (re.compile(r"enum (\w+)"), "enum"), (re.compile(r"trait (\w+)"), "trait"), (re.compile(r"impl (\w+)"), "impl"), (re.compile(r"func (\w+)"), "function"), (re.compile(r"type (\w+) struct"), "struct"), (re.compile(r"type (\w+) interface"), "interface"), (re.compile(r"function (\w+)"), "function"), (re.compile(r"async function (\w+)"), "function"), (re.compile(r"typedef struct (\w+)"), "typedef"), (re.compile(r"#define (\w+)"), "macro"), (re.compile(r"interface (\w+)"), "interface"), (re.compile(r"module (\w+)"), "module"), ] def extract_name(code: str) -> tuple[str, str]: for pat, typ in NAME_PATTERNS: m = pat.search(code) if m: return m.group(1), typ return "unknown", "block" def chunk_document(doc: dict) -> list[dict]: """Split a document into chunks based on language-specific patterns.""" content = doc["content"] lang = doc["language"] filepath = doc["filepath"] repo = doc["repo"] lines = content.split("\n") n_lines = len(lines) start_pat = CHUNK_STARTS.get(lang) if start_pat is None: name, typ = extract_name(content) return [{ "id": f"{repo}:{filepath}:0", "language": lang, "name": name, "type": typ, "code": content, "filepath": filepath, "repo": repo, "start_line": 1, "end_line": n_lines, }] starts = [(m.start(), m.group()) for m in start_pat.finditer(content)] if not starts: name, typ = extract_name(content) return [{ "id": f"{repo}:{filepath}:0", "language": lang, "name": name, "type": typ, "code": content, "filepath": filepath, "repo": repo, "start_line": 1, "end_line": n_lines, }] if starts[0][0] > 0: starts.insert(0, (0, "")) chunks = [] for i, (start_pos, _) in enumerate(starts): end_pos = starts[i + 1][0] if i + 1 < len(starts) else len(content) chunk_code = content[start_pos:end_pos].strip() if len(chunk_code) < 40: continue # Split huge chunks on blank lines if len(chunk_code) > 6000: sub_parts = re.split(r"\n\n+", chunk_code) for j, part in enumerate(sub_parts): part = part.strip() if len(part) < 40: continue name, typ = extract_name(part) start_line = content[:start_pos].count("\n") + 1 + sum(p.count("\n") + 2 for p in sub_parts[:j]) chunks.append({ "id": f"{repo}:{filepath}:{i}_{j}", "language": lang, "name": name, "type": typ, "code": part, "filepath": filepath, "repo": repo, "start_line": start_line, "end_line": start_line + part.count("\n"), }) continue name, typ = extract_name(chunk_code) start_line = content[:start_pos].count("\n") + 1 chunks.append({ "id": f"{repo}:{filepath}:{i}", "language": lang, "name": name, "type": typ, "code": chunk_code, "filepath": filepath, "repo": repo, "start_line": start_line, "end_line": start_line + chunk_code.count("\n"), }) return chunks def main(): os.makedirs(PRETRAIN_DIR, exist_ok=True) os.makedirs(CHUNKS_DIR, exist_ok=True) print(f"Loading quality documents from {QUALITY_DOCS_PATH}...") docs = [] with open(QUALITY_DOCS_PATH, "r", encoding="utf-8") as f: for line in f: docs.append(json.loads(line)) print(f" Loaded {len(docs):,} documents") print(f" Total size: {sum(len(d['content']) for d in docs) / 1e6:.1f} MB") # ─── Build pre-training corpus ───────────────────────────────────────── print("\nBuilding pre-training corpus...") corpus_size = 0 with open(CORPUS_PRETRAIN_PATH, "w", encoding="utf-8") as f: for doc in docs: f.write(doc["content"].strip()) f.write("\n\n\n") # document separator corpus_size += len(doc["content"]) print(f" Pre-training corpus: {corpus_size / 1e6:.1f} MB -> {CORPUS_PRETRAIN_PATH}") # ─── Build retrieval chunks ──────────────────────────────────────────── print("\nChunking documents...") all_chunks = [] lang_counts = {} type_counts = {} for doc in docs: chunks = chunk_document(doc) for chunk in chunks: all_chunks.append(chunk) lang_counts[chunk["language"]] = lang_counts.get(chunk["language"], 0) + 1 type_counts[chunk["type"]] = type_counts.get(chunk["type"], 0) + 1 print(f"\nTotal chunks: {len(all_chunks):,}") print(f" By language: {lang_counts}") print(f" By type: {type_counts}") sizes = [len(c["code"]) for c in all_chunks] if sizes: sizes.sort() print(f" Chunk size: min={sizes[0]}, median={sizes[len(sizes)//2]}, " f"max={sizes[-1]}, mean={sum(sizes)//len(sizes)}") with open(CHUNKS_PATH, "w", encoding="utf-8") as f: for chunk in all_chunks: f.write(json.dumps(chunk, ensure_ascii=False) + "\n") print(f"\nChunks written to {CHUNKS_PATH}") # Stats stats = { "input_docs": len(docs), "pretrain_corpus_mb": corpus_size / 1e6, "total_chunks": len(all_chunks), "language_distribution": lang_counts, "type_distribution": type_counts, } stats_path = os.path.join(CHUNKS_DIR, "chunk_stats.json") with open(stats_path, "w") as f: json.dump(stats, f, indent=2) print(f"Stats written to {stats_path}") if __name__ == "__main__": main()