Spaces:
Sleeping
Sleeping
File size: 3,255 Bytes
2f25a40 | 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 | """Ask a question against the indexed RL corpus.
Usage:
uv run python scripts/ask.py "What is the main idea of PPO?"
uv run python scripts/ask.py --k 8 "How does Rainbow combine Double DQN and PER?"
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel
from researchpath.embeddings import Embedder
from researchpath.index import load_index, search
from researchpath.rag import answer
from researchpath.retrieval import HybridRetriever, Reranker
console = Console()
ROOT = Path(__file__).resolve().parents[1]
INDEX_PATH = ROOT / "data" / "index.faiss"
def main() -> int:
parser = argparse.ArgumentParser(description="Ask the RL corpus a question.")
parser.add_argument("question", help="Question to ask.")
parser.add_argument("--k", type=int, default=5, help="Top-k chunks to retrieve (default: 5).")
parser.add_argument("--hybrid", action="store_true", help="Use BM25+FAISS hybrid retrieval.")
parser.add_argument("--rerank", action="store_true", help="Apply cross-encoder reranking on retrieved candidates.")
parser.add_argument(
"--show-sources",
action="store_true",
help="Print the retrieved chunks before the answer.",
)
args = parser.parse_args()
if not INDEX_PATH.exists():
console.print(f"[red]No index at {INDEX_PATH}. Run scripts/build_index.py first.[/red]")
return 1
if args.rerank:
mode = "hybrid+rerank (BM25+FAISS+CrossEncoder)" if args.hybrid else "dense+rerank (FAISS+CrossEncoder)"
else:
mode = "hybrid (BM25+FAISS)" if args.hybrid else "dense (FAISS)"
console.print(f"[bold cyan]Q:[/bold cyan] {args.question} [dim][retrieval: {mode}][/dim]\n")
index, chunks = load_index(INDEX_PATH)
embedder = Embedder()
if args.rerank:
base = HybridRetriever(index, chunks, embedder) if args.hybrid else None
if base is None:
# Wrap dense FAISS in a hybrid-compatible interface for the reranker
base = HybridRetriever(index, chunks, embedder)
reranker = Reranker(base)
hits = reranker.search(args.question, k=args.k)
elif args.hybrid:
retriever = HybridRetriever(index, chunks, embedder)
hits = retriever.search(args.question, k=args.k)
else:
hits = search(index, chunks, embedder, args.question, k=args.k)
if args.show_sources:
for i, h in enumerate(hits, 1):
console.print(
Panel(
h.text,
title=f"#{i} [{h.arxiv_id}, p{h.page}] score={h.score:.3f}",
border_style="dim",
)
)
console.print()
result = answer(args.question, hits)
console.print(Panel(Markdown(result.answer), title="Answer", border_style="green"))
console.print(
f"\n[dim]Retrieved {len(hits)} chunks | "
f"model: {result.llm.model} | "
f"tokens: {result.llm.input_tokens} in / {result.llm.output_tokens} out[/dim]"
)
return 0
if __name__ == "__main__":
sys.exit(main())
|