""" Chunk the curated corpus into searchable code units. Splits code into semantically meaningful chunks: - Python: functions, classes, top-level blocks - JS/TS: functions, classes, export blocks - Rust: fn, impl, struct, enum, trait blocks - Go: func, type, struct blocks - C/C++: function definitions, struct/typedef blocks Each chunk gets: - id: unique identifier - language: detected language - name: extracted name (function/class name) - type: function/class/struct/etc - code: the raw code text - filepath: synthetic path (derived from doc index) - start_line, end_line: line range within the doc Output: data/chunks.jsonl """ import json import os import re PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) CURATED_PATH = os.path.join(PROJECT_DIR, "data", "corpus_curated.txt") CHUNKS_PATH = os.path.join(PROJECT_DIR, "data", "chunks.jsonl") # ─── Language detection ────────────────────────────────────────────────────── LANG_PATTERNS = { "python": re.compile(r"^(def |class |import |from \S+ import |if __name__|@)", re.M), "js_ts": re.compile(r"^(function |const |let |class |export |import |async function|interface |type \w+ =)", re.M), "rust": re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |use |pub struct|pub enum|macro_rules!)", re.M), "go": re.compile(r"^(func |package |type \w+ struct|import )", re.M), "c_cpp": re.compile(r"^(#include|#define|#ifndef|#if |#endif|typedef |struct \w+|class \w+|void |int |char |size_t |ngx_)", re.M), } # ─── Chunk patterns per language ───────────────────────────────────────────── # Each pattern matches the START of a chunk. We split on these. CHUNK_STARTS = { "python": re.compile(r"^(def |class |@|if __name__)", re.M), "js_ts": re.compile(r"^(function |const |let |class |export |async function|interface |type \w+ =)", re.M), "rust": re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |macro_rules!|pub struct|pub enum)", re.M), "go": re.compile(r"^(func |type \w+ struct|type \w+ interface)", re.M), "c_cpp": re.compile(r"^(static |void |int |char |size_t |ngx_|typedef |struct \w+ \{|#define|#if|#ifndef)", re.M), } # ─── Name extraction patterns ──────────────────────────────────────────────── 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"function (\w+)"), "function"), (re.compile(r"typedef struct (\w+)"), "typedef"), (re.compile(r"#define (\w+)"), "macro"), (re.compile(r"(ngx_\w+)\s*\("), "function"), ] def detect_language(doc: str) -> str | None: for lang, pat in LANG_PATTERNS.items(): if len(pat.findall(doc)) >= 2: return lang return None def extract_name(code: str) -> tuple[str, str]: """Extract the name and type from a code chunk.""" for pat, typ in NAME_PATTERNS: m = pat.search(code) if m: return m.group(1), typ return "unknown", "block" def chunk_document(doc: str, lang: str, doc_idx: int) -> list[dict]: """Split a document into chunks based on language-specific patterns.""" lines = doc.split("\n") n_lines = len(lines) # Find all chunk start positions start_pat = CHUNK_STARTS.get(lang) if start_pat is None: # Fallback: treat whole doc as one chunk return [{ "id": f"doc_{doc_idx}_chunk_0", "language": lang, "name": "block", "type": "block", "code": doc, "filepath": f"src/doc_{doc_idx}.txt", "start_line": 1, "end_line": n_lines, }] starts = [(m.start(), m.group()) for m in start_pat.finditer(doc)] if not starts: # No pattern matches — treat whole doc as one chunk return [{ "id": f"doc_{doc_idx}_chunk_0", "language": lang, "name": "block", "type": "block", "code": doc, "filepath": f"src/doc_{doc_idx}.txt", "start_line": 1, "end_line": n_lines, }] # Add doc start if first match isn't at position 0 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(doc) chunk_code = doc[start_pos:end_pos].strip() # Skip tiny chunks (< 30 chars) if len(chunk_code) < 30: continue # Skip huge chunks (> 8000 chars — split them) if len(chunk_code) > 8000: # Split on blank lines sub_parts = re.split(r"\n\n+", chunk_code) for j, part in enumerate(sub_parts): if len(part.strip()) < 30: continue name, typ = extract_name(part) start_line = doc[:start_pos].count("\n") + 1 + sum(p.count("\n") + 2 for p in sub_parts[:j]) chunks.append({ "id": f"doc_{doc_idx}_chunk_{i}_{j}", "language": lang, "name": name, "type": typ, "code": part.strip(), "filepath": f"src/doc_{doc_idx}.txt", "start_line": start_line, "end_line": start_line + part.count("\n"), }) continue name, typ = extract_name(chunk_code) start_line = doc[:start_pos].count("\n") + 1 chunks.append({ "id": f"doc_{doc_idx}_chunk_{i}", "language": lang, "name": name, "type": typ, "code": chunk_code, "filepath": f"src/doc_{doc_idx}.txt", "start_line": start_line, "end_line": start_line + chunk_code.count("\n"), }) return chunks def main(): print(f"Loading curated corpus from {CURATED_PATH}...") with open(CURATED_PATH, "r", encoding="utf-8") as f: text = f.read() print(f" Corpus size: {len(text) / 1e6:.2f} MB") # Split into documents docs = re.split(r"\n{3,}", text) print(f" Documents: {len(docs):,}") # Chunk each document print("Chunking documents...") all_chunks = [] lang_counts = {} type_counts = {} for i, doc in enumerate(docs): doc = doc.strip() if len(doc) < 50: continue lang = detect_language(doc) if lang is None: continue chunks = chunk_document(doc, lang, i) for chunk in chunks: all_chunks.append(chunk) lang_counts[lang] = lang_counts.get(lang, 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}") # Size distribution sizes = [len(c["code"]) for c in all_chunks] sizes.sort() print(f" Chunk size: min={sizes[0]}, median={sizes[len(sizes)//2]}, max={sizes[-1]}, mean={sum(sizes)//len(sizes)}") # Write chunks with open(CHUNKS_PATH, "w", encoding="utf-8") as f: for chunk in all_chunks: f.write(json.dumps(chunk) + "\n") print(f"\nChunks written to {CHUNKS_PATH}") # Show a few samples print("\n" + "=" * 60) print("SAMPLE CHUNKS") print("=" * 60) for i in [0, len(all_chunks) // 4, len(all_chunks) // 2, len(all_chunks) - 1]: c = all_chunks[i] print(f"\n--- {c['id']} | {c['language']} | {c['type']} | {c['name']} ---") print(f" File: {c['filepath']}:{c['start_line']}-{c['end_line']}") print(f" Size: {len(c['code'])} chars") print(c["code"][:300]) if len(c["code"]) > 300: print("...") if __name__ == "__main__": main()