| """
|
| 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_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
|
|
|
|
|
| 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")
|
|
|
|
|
| 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")
|
| corpus_size += len(doc["content"])
|
| print(f" Pre-training corpus: {corpus_size / 1e6:.1f} MB -> {CORPUS_PRETRAIN_PATH}")
|
|
|
|
|
| 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 = {
|
| "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()
|
|
|
