ROS 2 docs RAG index and tooling

scripts/ask.py

@@ -0,0 +1,127 @@
+"""Answer a ROS 2 question from the local index.
+
+Retrieval and generation are deliberately separable. `retrieve()` needs only
+FAISS and the embedding model — a few hundred MB — and is useful on its own for
+"where is this documented?". Generation loads Qwen on top of that.
+
+    python scripts/ask.py "How do I define a devicetree binding?"
+    python scripts/ask.py "What is a work queue?" --k 8
+    python scripts/ask.py "How do I enable logging?" --retrieve-only
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+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")
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import retrieval  # noqa: E402
+
+ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_INDEX = ROOT / "data" / "index"
+DEFAULT_MODEL = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
+
+SYSTEM_PROMPT = """You are a Zephyr RTOS assistant. Answer only from the documentation \
+excerpts provided.
+
+Zephyr is large and moves fast, and a confident wrong answer costs an engineer \
+hours on hardware. If the excerpts do not cover the question, say so and name \
+what is missing rather than filling the gap from general RTOS knowledge.
+
+Cite the source path for anything you assert. Prefer Zephyr's own terminology \
+(devicetree, Kconfig, west, k_work) over generic equivalents. Show configuration \
+as it would actually appear — a Kconfig symbol, a devicetree node, a west \
+command — not as prose describing it."""
+
+
+def retrieve(question: str, k: int = 6, index_dir: Path = DEFAULT_INDEX) -> list[dict]:
+    """Shared with the Space via scripts/retrieval.py, including the
+    release-note down-weight — a CLI that ranked differently from the hosted
+    app would make every comparison between them meaningless."""
+    return retrieval.search(question, k, index_dir)
+
+
+def build_context(hits: list[dict]) -> str:
+    return "\n\n".join(
+        f"[{index + 1}] {hit['source']} - {hit['section']}\n{hit['text']}"
+        for index, hit in enumerate(hits)
+    )
+
+
+def generate(question: str, hits: list[dict], model_id: str, max_new_tokens: int = 512) -> str:
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+    )
+
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": f"Documentation excerpts:\n\n{build_context(hits)}\n\nQuestion: {question}",
+        },
+    ]
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    output = model.generate(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        do_sample=False,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+    return tokenizer.decode(output[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True).strip()
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("question")
+    parser.add_argument("--k", type=int, default=6)
+    parser.add_argument("--index", type=Path, default=DEFAULT_INDEX)
+    parser.add_argument("--model", default=DEFAULT_MODEL)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument(
+        "--retrieve-only",
+        action="store_true",
+        help="show the retrieved passages and stop - no model download or load",
+    )
+    args = parser.parse_args()
+
+    hits = retrieve(args.question, args.k, args.index)
+    if not hits:
+        print("Nothing retrieved.")
+        return 1
+
+    print(f"\nRetrieved {len(hits)} passages:\n")
+    for position, hit in enumerate(hits, start=1):
+        print(f"  [{position}] {hit['score']:.3f}  {hit['source']} - {hit['section']}")
+
+    if args.retrieve_only:
+        print()
+        for position, hit in enumerate(hits, start=1):
+            print(f"--- [{position}] {hit['source']} ---")
+            print(hit["text"][:500])
+            print()
+        return 0
+
+    print(f"\nGenerating with {args.model}...\n")
+    print(generate(args.question, hits, args.model, args.max_new_tokens))
+    print("\nSources:")
+    for hit in hits:
+        print(f"  {hit['source']} -> {hit['url']}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())