"""Chunk the Zephyr docs and build a FAISS index over them. Chunking is the part that decides whether retrieval works, so it is done on document structure rather than on a fixed character count. RST carries its own section headings; splitting on those keeps a chunk to one topic and lets every retrieved passage cite the section it came from. A blind 1000-character split would cut mid-sentence and mix two subsystems into one vector. python scripts/build_index.py python scripts/build_index.py --docs data/raw_docs --out data/index """ from __future__ import annotations import argparse import json import re import sys from pathlib import Path if hasattr(sys.stdout, "reconfigure"): sys.stdout.reconfigure(encoding="utf-8", errors="replace") sys.stderr.reconfigure(encoding="utf-8", errors="replace") ROOT = Path(__file__).resolve().parent.parent DEFAULT_DOCS = ROOT / "data" / "raw_docs" DEFAULT_OUT = ROOT / "data" / "index" EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2" # An RST section underline: a run of one punctuation character on its own line, # directly under the title it underlines. RST_UNDERLINE = re.compile(r"^([=\-`:'\"~^_*+#<>])\1{2,}\s*$") MIN_CHARS = 120 MAX_CHARS = 1600 # Files that live under doc/ but are not documentation. Indexing them puts a # pip requirements file and a JavaScript 404 page into the retrieval pool. SKIP_FILES = { "404.rst", "CMakeLists.txt", "requirements.txt", "substitutions.txt", "index.rst", } # Release notes end with thousands of bare GitHub issue IDs. They are the # largest sections in the corpus and answer nothing: one release note produced a # single 104 KB chunk of numbers. SKIP_SECTIONS = re.compile(r"issue related items|^bugs?$|security vulnerability", re.IGNORECASE) # A chunk that is mostly punctuation, IDs or markup is not prose. Retrieval will # happily return it and the model will have nothing to say about it. def looks_like_prose(text: str) -> bool: letters = sum(character.isalpha() for character in text) if letters < len(text) * 0.45: return False words = text.split() if not words: return False # Long runs of short tokens are ID lists and option tables, not sentences. return sum(len(word) for word in words) / len(words) >= 3.0 def clean(text: str) -> str: """Strip the RST directives that carry no meaning for a reader.""" lines: list[str] = [] for line in text.splitlines(): stripped = line.strip() # Comments, toctrees and figure/image directives are navigation and # layout. They retrieve badly and answer nothing. if stripped.startswith(".. toctree::") or stripped.startswith(".. figure::"): continue if stripped.startswith(".. image::") or stripped.startswith(".."): if re.match(r"^\.\.\s+_[\w.-]+:", stripped): continue # anchor target if stripped.startswith(".. code-block::") or stripped.startswith(".. note::"): lines.append(line) # keep — the body that follows is content continue continue if stripped.startswith(":") and stripped.count(":") >= 2 and len(stripped) < 80: continue # field list / directive option lines.append(line) return "\n".join(lines) def sections(text: str) -> list[tuple[str, str]]: """Split RST into (heading, body). The first block inherits the document title.""" lines = text.splitlines() blocks: list[tuple[str, list[str]]] = [("", [])] index = 0 while index < len(lines): line = lines[index] following = lines[index + 1] if index + 1 < len(lines) else "" is_heading = ( line.strip() and RST_UNDERLINE.match(following) and len(following.strip()) >= len(line.strip()) - 2 ) if is_heading: blocks.append((line.strip(), [])) index += 2 continue blocks[-1][1].append(line) index += 1 return [(heading, "\n".join(body).strip()) for heading, body in blocks] def split_long(body: str, limit: int = MAX_CHARS) -> list[str]: """Break an over-long section on blank lines, never mid-paragraph. Paragraph splitting alone is not enough. A section with no blank lines — a long table, a generated list — comes back as one piece however large it is. The first version of this used only blank lines and emitted a single 126 KB chunk, so anything still over the limit is hard-split on line boundaries as a backstop. """ if len(body) <= limit: return [body] parts: list[str] = [] current = "" for paragraph in body.split("\n\n"): if current and len(current) + len(paragraph) + 2 > limit: parts.append(current.strip()) current = paragraph else: current = f"{current}\n\n{paragraph}" if current else paragraph if current.strip(): parts.append(current.strip()) bounded: list[str] = [] for part in parts: if len(part) <= limit: bounded.append(part) continue buffer = "" for line in part.splitlines(): if buffer and len(buffer) + len(line) + 1 > limit: bounded.append(buffer.strip()) buffer = line else: buffer = f"{buffer}\n{line}" if buffer else line if buffer.strip(): bounded.append(buffer.strip()) return bounded def chunk_document(path: Path) -> list[dict]: raw = path.read_text(encoding="utf-8", errors="replace") # The fetcher flattens "kernel/services/threads.rst" to # "kernel__services__threads.rst" so the source path survives as a citation. source = path.stem.replace("__", "/") body = clean(raw) title = "" for heading, _ in sections(body): if heading: title = heading break chunks: list[dict] = [] for heading, section_body in sections(body): if len(section_body) < MIN_CHARS: continue if heading and SKIP_SECTIONS.search(heading): continue for part in split_long(section_body): if len(part) < MIN_CHARS or not looks_like_prose(part): continue chunks.append( { "text": part, "source": source, "title": title or source, "section": heading or title or source, "url": f"https://docs.zephyrproject.org/latest/{source}.html", } ) return chunks def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--docs", type=Path, default=DEFAULT_DOCS) parser.add_argument("--out", type=Path, default=DEFAULT_OUT) parser.add_argument("--model", default=EMBEDDING_MODEL) parser.add_argument("--batch-size", type=int, default=64) args = parser.parse_args() if not args.docs.is_dir(): raise SystemExit(f"no docs at {args.docs} — run scripts/fetch_docs.py first") files = sorted(p for p in args.docs.glob("*") if p.suffix.lower() in {".rst", ".md", ".txt"}) files = [p for p in files if p.name != "_SOURCE.txt" and p.name not in SKIP_FILES] print(f"Chunking {len(files)} documents...") chunks: list[dict] = [] for path in files: chunks.extend(chunk_document(path)) if not chunks: raise SystemExit("no chunks produced — check the docs directory") lengths = sorted(len(c["text"]) for c in chunks) print( f"{len(chunks)} chunks | median {lengths[len(lengths) // 2]} chars, " f"max {lengths[-1]}" ) # Imported here so --help and the chunking stats stay fast on a machine # without torch installed. import faiss import numpy as np from sentence_transformers import SentenceTransformer print(f"Embedding with {args.model}...") model = SentenceTransformer(args.model) vectors = model.encode( [c["text"] for c in chunks], batch_size=args.batch_size, show_progress_bar=True, convert_to_numpy=True, normalize_embeddings=True, ).astype("float32") # Inner product over L2-normalised vectors is cosine similarity, which is # what the retrieval scores in ask.py are reported as. index = faiss.IndexFlatIP(vectors.shape[1]) index.add(vectors) args.out.mkdir(parents=True, exist_ok=True) faiss.write_index(index, str(args.out / "docs.faiss")) with (args.out / "chunks.jsonl").open("w", encoding="utf-8") as handle: for chunk in chunks: handle.write(json.dumps(chunk, ensure_ascii=False) + "\n") source_note = (args.docs / "_SOURCE.txt") meta = { "embedding_model": args.model, "dimensions": int(vectors.shape[1]), "chunks": len(chunks), "documents": len(files), "index": "IndexFlatIP (cosine over normalised vectors)", "source": source_note.read_text(encoding="utf-8") if source_note.exists() else "unknown", } (args.out / "meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8") print(f"\nIndex written to {args.out}") print(f"{len(chunks)} chunks x {vectors.shape[1]} dims") return 0 if __name__ == "__main__": sys.exit(main())