scripts/ask.py

"""Answer a Zephyr 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())