ThinkRetrieve in 10 minutes

This tutorial has two starting tracks: an offline smoke test requiring no model, dataset, GPU, or API key, and a real run using a local model or hosted API. All repository commands below run from the repository root--the directory containing README.md and thinkretrieve/.

0. Create an environment

git clone https://github.com/itsvaibhav01/ThinkRetrieve.git
cd ThinkRetrieve

python -m venv .venv
source .venv/bin/activate              # Windows: .venv\Scripts\activate
python -m pip install --upgrade pip

1. Offline smoke test--no model or dataset downloads

python -m pip install -e ./thinkretrieve
python thinkretrieve/examples/offline_dummy.py

This deterministic example exercises the real think -> retrieve -> inject -> continue loop and should print Final answer: \$312.50 and Retrievals: 1. It verifies installation and control flow; it is not an accuracy benchmark.

2. Prepare a real local model

Install Ollama, then download the tutorial model:

ollama pull qwen3:4b

Start the Ollama application (or run ollama serve), then install retrieval and dataset dependencies:

python -m pip install -e "./thinkretrieve[faiss]" datasets

The first retrieval run also downloads an E5 encoder through SentenceTransformers. qwen3:4b is only the tutorial default: you may use a smaller or larger Ollama model, another OpenAI-compatible server/API, Anthropic, or Amazon Bedrock.

---

A. See the gain yourself: index -> infer -> compare

The comparison downloads GSM8K and intfloat/e5-base-v2, indexes 4,000 training solutions (usually a few minutes on CPU after downloads), then evaluates plain querying, prompt-level RAG, long thinking, and ThinkRetrieve at the same budget. It writes gsm8k_bank/ and resumable ab_results_*.json checkpoints in the current directory.

Start with a small pipeline check. It is too small for accuracy conclusions:

python thinkretrieve/examples/compare_tts_vs_thinkretrieve.py \
    --backend openai --model qwen3:4b \
    --base-url http://localhost:11434/v1 \
    --limit 2 --bank-size 100 --budget 512

Then run the full defaults:

Amazon Bedrock (needs pip install boto3 + AWS credentials):

python thinkretrieve/examples/compare_tts_vs_thinkretrieve.py \
    --backend bedrock --model qwen.qwen3-32b-v1:0 --region us-east-1

Local model on a Mac (needs Ollama: ollama pull qwen3:4b):

python thinkretrieve/examples/compare_tts_vs_thinkretrieve.py \
    --backend openai --model qwen3:4b --base-url http://localhost:11434/v1

Output (shape — your numbers depend on model/budget):

                        accuracy  thinking tok/q  retrievals/q
long thinking                 x%            ~2k             —
thinkretrieve                 y%            ~2k           1.x
==========================================================
Δ accuracy: +z% on 25 GSM8K questions (same model, same thinking budget)

Every question is checkpointed to ab_results_*.json, so you can interrupt and resume, raise --limit, or swap --budget. To feel the effect where it's biggest, use a small model (1–4B) and a hard set — in the paper, Qwen3-1.7B goes 22.2% → 35.6% on AIME 2025 and the method wins every (model, benchmark) cell across GSM-8K, MATH-500, AIME 2025, and SciQ.

What the script actually does (the whole API)

from thinkretrieve import ThinkRetrieve, FaissRetriever, OpenAICompatBackend

# 1) INDEX — any (question, worked-solution) pairs
bank = FaissRetriever.from_examples(
    [(row["question"], row["answer"]) for row in gsm8k_train]
)
bank.save("gsm8k_bank")                     # reuse forever

# 2) INFERENCE — think → try to answer → retrieve similar solved example
#    → inject it into the thinking trace → continue → answer
tr = ThinkRetrieve(backend, bank)
result = tr.run("A robe takes 2 bolts of blue fiber and half that much white...")

result.answer              # final answer
result.think_trace         # full reasoning, injected examples visible
result.retrievals          # which examples fired, when, and why

Baseline for comparison is one argument: ThinkRetrieve(backend, retriever=None) is exactly the "think longer" baseline (budget forcing) — same loop, no retrieval — so the A/B is always apples-to-apples.

---

B. Procedural memory for agents

Most agent memory stores facts. A ThinkRetrieve bank stores procedures — solved tasks with their worked solutions — and recalls them mid-reasoning, at the moment the model is about to commit to an answer:

from thinkretrieve import FaissRetriever, ThinkRetrieve

memory = FaissRetriever.load("agent_procedural_memory")
agent  = ThinkRetrieve(backend, memory)

result = agent.run(task)

if verified(result):                 # your check: tests pass, human approves...
    memory.add(task, result.answer)  # the agent now remembers HOW
    memory.save("agent_procedural_memory")

The bank grows as the agent works; tomorrow's similar task retrieves today's verified solution. See examples/agent_memory.py for a runnable version:

ollama pull qwen3:4b
python thinkretrieve/examples/agent_memory.py

---

C. Already have a RAG system? Change where you inject

The paper's most practical ablation: prepending retrieved content to the prompt (standard RAG) was ≈ neutral on reasoning tasks (88.8% vs 89.1% baseline), while injecting the same retrieved content inside the thinking trace produced the gain. You don't need a new index — wrap the one you have:

from thinkretrieve.types import Example

class MyVectorDB:                     # anything: pgvector, Pinecone, BM25...
    def retrieve(self, question, interim_answer="", k=1, exclude_ids=frozenset()):
        hits = my_index.search(f"{question}\n{interim_answer}", k + len(exclude_ids))
        return [Example(id=h.id, question=h.title, solution=h.body)
                for h in hits if h.id not in exclude_ids][:k]

tr = ThinkRetrieve(backend, MyVectorDB())

Two things you get for free that prompt-level RAG can't do: the query includes the model's interim answer (its current reasoning state, not just the question), and injection happens at each reasoning boundary with per-question deduplication.

---

Other runnable paths

Tiny local quickstart

After ollama pull qwen3:4b, run:

python thinkretrieve/examples/quickstart.py

This creates a small in-memory bank. On its first run, SentenceTransformers downloads intfloat/e5-large-v2.

SciQ

The SciQ example downloads the SciQ train/test splits and intfloat/e5-base-v2, builds a bank, and queries the configured model:

python thinkretrieve/examples/sciq_example.py \
    --model qwen3:4b --limit 2 --bank-size 100

Increase to --limit 30 --bank-size 3000 after the smoke run.

NuminaMath bank

python thinkretrieve/examples/build_bank_from_numina.py

This is a full data-building job: it downloads NuminaMath-1.5, selects 20,000 synthetic examples, embeds them, and writes numina_bank/. Expect a substantial download, CPU time, and disk use. It is not required for the offline or tiny quickstarts.

Anthropic

python -m pip install -e "./thinkretrieve[faiss,anthropic]"
export ANTHROPIC_API_KEY=YOUR_KEY
python thinkretrieve/examples/quickstart_anthropic.py

The Anthropic example expects a saved bank at my_bank/; create it first as shown in quickstart_local_mac.py. API usage may incur provider charges.

Hosted OpenAI-compatible APIs

Use the comparison CLI's --base-url, --model, and --api-key arguments, or configure OpenAICompatBackend directly. Never commit keys or .env files.

---

Choosing knobs

Knob	Default	Advice
think_budget	8192	The headroom for gains; paper sweeps 2K–22K.
min_seg	1024	Bigger = fewer API round-trips = cheaper without prefix caching.
max_insertions	8	2–4 is plenty at small budgets.
num_examples	1	Paper default; more examples eat budget fast.
encoder	e5-large-v2	e5-base-v2 is 3× faster and nearly as good (paper §5: gains are encoder-insensitive).

Backends: OpenAICompatBackend (Ollama, LM Studio, MLX, llama.cpp, vLLM, OpenAI, Together, Groq, DeepSeek, OpenRouter), AnthropicBackend, BedrockConverseBackend — or subclass ChatTranscriptBackend and implement one _chat() method.

Ollama 404 / model-not-found tip: run ollama pull qwen3:4b, check ollama list, and make sure --model exactly matches an installed name. If the connection is refused instead, start the app or run ollama serve.

TensorFlow/protobuf tip: SentenceTransformers uses PyTorch here. If an unrelated TensorFlow installation has incompatible protobuf packages, run the example with USE_TF=0.

macOS segfault tip: if a script that mixes FAISS with torch or datasets/pyarrow crashes with a bare Segmentation fault on Apple Silicon, it's a native OpenMP runtime clash in faiss.search(). Run with OMP_NUM_THREADS=1 (single-threaded search is more than fast enough for per-question retrieval), and prefer FaissRetriever.load(..., device="cpu") for the query encoder at inference time.

Qwen3 tip: Qwen3 models re-open a <think> block on every pass, including the short interim/final answer passes, which can eat the answer budget. Qwen3 honors a /no_think soft switch — append it to the answer instructions:

from thinkretrieve.backends import base

class Qwen3Backend(OpenAICompatBackend):
    final_answer_instruction = base.FINAL_ANSWER_INSTRUCTION + " /no_think"