scripts/pilot_retrieval.py

#!/usr/bin/env python3
"""Build and inspect a small, deterministic MLX semantic-retrieval pilot."""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from mlx_lm import load


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_INDEX_DIR = PROJECT_ROOT / "data" / "index" / "pilot"
MODEL_ID = "mlx-community/Qwen3-Embedding-0.6B-4bit-DWQ"
MODEL_REVISION = "6c3ae70858513f1a78e9cdca3cae330d9075cd2a"
TASK_INSTRUCTION = (
    "Given a control-systems engineering query, retrieve technically relevant "
    "passages, equations, and executable MATLAB or Python examples."
)

CHUNK_FILES = [
    PROJECT_ROOT / "data" / "processed" / "chunks" / "knowledge_chunks.jsonl",
    PROJECT_ROOT
    / "data"
    / "processed"
    / "web_collections_chunks"
    / "knowledge_chunks.jsonl",
    PROJECT_ROOT
    / "data"
    / "processed"
    / "core_books_chunks"
    / "knowledge_chunks.jsonl",
    PROJECT_ROOT
    / "data"
    / "processed"
    / "arxiv_chunks"
    / "knowledge_chunks.jsonl",
]

TOPICS = {
    "controllability": [
        "controllability",
        "controllability matrix",
        "kalman rank",
        "reachable subspace",
        "ctrb",
    ],
    "h_infinity": [
        "h infinity",
        "hinfinity",
        "mixed sensitivity",
        "weighting function",
        "small gain",
        "hinfsyn",
    ],
    "mpc": [
        "model predictive control",
        "receding horizon",
        "finite horizon",
        "input constraints",
        "state constraints",
        "terminal cost",
    ],
}

DEMO_QUERIES = [
    "How do I test controllability of a continuous-time LTI system and compute the controllability matrix?",
    "Explain mixed-sensitivity H-infinity synthesis using weighting functions on S, KS, and T.",
    "How does model predictive control enforce input and state constraints over a finite horizon?",
]


def load_chunks() -> list[dict]:
    chunks = []
    for path in CHUNK_FILES:
        if not path.exists():
            continue
        with path.open(encoding="utf-8") as stream:
            for line in stream:
                row = json.loads(line)
                if row.get("text"):
                    row["chunk_file"] = str(path.relative_to(PROJECT_ROOT))
                    chunks.append(row)
    return chunks


def keyword_score(text: str, terms: list[str]) -> int:
    lowered = text.lower().replace("-", " ")
    return sum(lowered.count(term.replace("-", " ")) for term in terms)


def select_pilot_chunks(chunks: list[dict], per_topic: int, distractors: int) -> list[dict]:
    selected: dict[str, dict] = {}
    for topic, terms in TOPICS.items():
        ranked = sorted(
            chunks,
            key=lambda row: (
                keyword_score(row["text"], terms),
                row["token_count"],
                row["chunk_id"],
            ),
            reverse=True,
        )
        for row in (candidate for candidate in ranked if keyword_score(candidate["text"], terms) > 0):
            copy = dict(row)
            copy["pilot_topic"] = topic
            selected.setdefault(copy["chunk_id"], copy)
            if sum(item.get("pilot_topic") == topic for item in selected.values()) >= per_topic:
                break

    for row in sorted(chunks, key=lambda item: item["chunk_id"]):
        if row["chunk_id"] in selected:
            continue
        copy = dict(row)
        copy["pilot_topic"] = "distractor"
        selected[copy["chunk_id"]] = copy
        distractors -= 1
        if distractors == 0:
            break
    return list(selected.values())


def query_text(query: str) -> str:
    return f"Instruct: {TASK_INSTRUCTION}\nQuery:{query}"


def embed_text(model, tokenizer, text: str, max_tokens: int) -> np.ndarray:
    token_ids = tokenizer.encode(text, add_special_tokens=True)
    token_ids = token_ids[-max_tokens:]
    hidden = model.model(mx.array([token_ids]))
    vector = hidden[0, -1].astype(mx.float32)
    vector = vector / mx.sqrt(mx.sum(vector * vector))
    mx.eval(vector)
    return np.asarray(vector, dtype=np.float32)


def build_index(index_dir: Path, per_topic: int, distractors: int, max_tokens: int) -> None:
    chunks = select_pilot_chunks(load_chunks(), per_topic, distractors)
    if not chunks:
        raise RuntimeError("No chunks were found. Build the processed corpus first.")

    print(f"Loading {MODEL_ID} at {MODEL_REVISION[:8]}...")
    model, tokenizer = load(MODEL_ID, revision=MODEL_REVISION)
    mx.reset_peak_memory()
    started = time.perf_counter()
    vectors = []
    for index, chunk in enumerate(chunks, start=1):
        vectors.append(embed_text(model, tokenizer, chunk["text"], max_tokens))
        if index % 10 == 0 or index == len(chunks):
            print(f"Embedded {index}/{len(chunks)} chunks")
    elapsed = time.perf_counter() - started

    matrix = np.stack(vectors).astype(np.float16)
    index_dir.mkdir(parents=True, exist_ok=True)
    np.save(index_dir / "embeddings.npy", matrix)
    with (index_dir / "metadata.jsonl").open("w", encoding="utf-8") as stream:
        for chunk in chunks:
            stream.write(json.dumps(chunk, ensure_ascii=False) + "\n")
    manifest = {
        "model_id": MODEL_ID,
        "model_revision": MODEL_REVISION,
        "task_instruction": TASK_INSTRUCTION,
        "chunks": len(chunks),
        "dimensions": int(matrix.shape[1]),
        "dtype": str(matrix.dtype),
        "max_tokens": max_tokens,
        "elapsed_seconds": elapsed,
        "chunks_per_second": len(chunks) / elapsed,
        "peak_mlx_memory_bytes": mx.get_peak_memory(),
    }
    (index_dir / "manifest.json").write_text(
        json.dumps(manifest, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(manifest, indent=2))


def search(index_dir: Path, queries: list[str], top_k: int, max_tokens: int) -> None:
    vectors = np.load(index_dir / "embeddings.npy").astype(np.float32)
    with (index_dir / "metadata.jsonl").open(encoding="utf-8") as stream:
        chunks = [json.loads(line) for line in stream]

    print(f"Loading {MODEL_ID} for query embedding...")
    model, tokenizer = load(MODEL_ID, revision=MODEL_REVISION)
    for query in queries:
        query_vector = embed_text(model, tokenizer, query_text(query), max_tokens)
        scores = vectors @ query_vector
        indices = np.argsort(scores)[::-1][:top_k]
        print(f"\nQUERY: {query}")
        for rank, index in enumerate(indices, start=1):
            chunk = chunks[int(index)]
            excerpt = " ".join(chunk["text"].split())[:260]
            print(
                f"{rank}. score={scores[index]:.4f} topic={chunk['pilot_topic']} "
                f"source={chunk['source_id']}\n   {excerpt}"
            )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--index-dir", type=Path, default=DEFAULT_INDEX_DIR)
    parser.add_argument("--build", action="store_true")
    parser.add_argument("--query", action="append", default=[])
    parser.add_argument("--top-k", type=int, default=5)
    parser.add_argument("--per-topic", type=int, default=15)
    parser.add_argument("--distractors", type=int, default=15)
    parser.add_argument("--max-tokens", type=int, default=1024)
    args = parser.parse_args()

    if args.build:
        build_index(args.index_dir, args.per_topic, args.distractors, args.max_tokens)
    queries = args.query or DEMO_QUERIES
    search(args.index_dir, queries, args.top_k, args.max_tokens)


if __name__ == "__main__":
    main()