File size: 8,762 Bytes
3738348 | 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 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 | """
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()
|