File size: 3,421 Bytes
c7815d1 | 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 | """Shared retrieval for the CLI and the Space, so the two cannot drift apart.
Beyond loading the index, this does one thing worth explaining: it down-weights
release notes.
Zephyr's release notes are a quarter of the corpus — 1,688 chunks of the 6,949 —
and they are written in the same vocabulary as the documentation they describe.
A question like "what does west build actually do?" retrieved
`develop/west/release-notes - v0.7.0` above `develop/west/build-flash-debug`,
because a changelog entry about west build looks a lot like documentation about
west build to an embedding model.
For a coding assistant that is the wrong answer nearly every time: a developer
asking how something works wants the reference page, not what changed in v0.7.0.
So release notes are penalised, unless the question is actually about a version
or a change — in which case they are exactly right and the penalty is lifted.
"""
from __future__ import annotations
import json
import re
from functools import lru_cache
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
DEFAULT_INDEX = ROOT / "data" / "index"
# Historical rather than instructional. Not excluded — "what changed in 3.5?" is
# a fair question — just made to earn its place.
HISTORICAL = re.compile(r"^releases/|/release-notes", re.IGNORECASE)
HISTORICAL_PENALTY = 0.08
# When the question is about a version or a change, the penalty is lifted.
VERSION_QUERY = re.compile(
r"\b(v?\d+\.\d+|release|changelog|changed|new in|migration|deprecat\w*|"
r"breaking|since|upgrade|version)\b",
re.IGNORECASE,
)
@lru_cache(maxsize=4)
def load(index_dir: str):
import faiss
from sentence_transformers import SentenceTransformer
path = Path(index_dir)
if not (path / "docs.faiss").exists():
raise FileNotFoundError(
f"no index at {path} - run scripts/fetch_docs.py then scripts/build_index.py"
)
meta = json.loads((path / "meta.json").read_text(encoding="utf-8"))
chunks = [
json.loads(line)
for line in (path / "chunks.jsonl").read_text(encoding="utf-8").splitlines()
]
embedder = SentenceTransformer(meta["embedding_model"])
return faiss.read_index(str(path / "docs.faiss")), chunks, embedder, meta
def search(question: str, k: int = 6, index_dir: Path | str = DEFAULT_INDEX) -> list[dict]:
index, chunks, embedder, _ = load(str(index_dir))
vector = embedder.encode(
[question], convert_to_numpy=True, normalize_embeddings=True
).astype("float32")
# Over-fetch so re-ranking has something to promote from. Without this the
# penalty could only demote within an already-decided top k.
scores, ids = index.search(vector, min(k * 4, len(chunks)))
wants_history = bool(VERSION_QUERY.search(question))
hits: list[dict] = []
for score, position in zip(scores[0], ids[0]):
if position < 0:
continue
hit = dict(chunks[int(position)])
raw = float(score)
adjusted = raw
if not wants_history and HISTORICAL.search(hit.get("source", "")):
adjusted -= HISTORICAL_PENALTY
hit["score"] = round(adjusted, 4)
hit["raw_score"] = round(raw, 4)
hits.append(hit)
hits.sort(key=lambda item: item["score"], reverse=True)
return hits[:k]
def meta(index_dir: Path | str = DEFAULT_INDEX) -> dict:
return load(str(index_dir))[3]
|