File size: 4,661 Bytes
4a0e2e8 | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """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())
|