ThinkRetrieve

September 5, 2026 · View on GitHub

ThinkRetrieve

PyPI Python License arXiv

Paper · PDF · Project page · Abstract & citations · Library · Tutorial · Article · Science-agent case study


⭐ Accepted for EMNLP 2026 findings

ThinkRetrieve

ThinkRetrieve is retrieval-augmented test-time scaling for reasoning models. While a model reasons, it retrieves a similar worked example and injects that example into the live thinking trace. The model receives guidance on how to reason, not only additional facts in its initial prompt.

This repository contains two public-facing parts:

ThinkRetrieve pipeline

Choose your starting point

1. Fastest check: no model, dataset, API key, or GPU

This verifies the complete think → retrieve → inject → continue loop with a small deterministic backend. It does not measure model quality.

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 -e ./thinkretrieve

python thinkretrieve/examples/offline_dummy.py

2. Run with a local model through Ollama

Install Ollama, then download the tutorial model:

ollama pull qwen3:4b

Start the Ollama application, or run ollama serve, and then:

python -m pip install -e "./thinkretrieve[faiss]"
python thinkretrieve/examples/quickstart.py

The first run also downloads the E5 embedding model used to build the example bank. Generation runs in Ollama; retrieval runs locally on CPU. qwen3:4b is the tutorial default, not a requirement—you can use a smaller or larger local model, an OpenAI-compatible hosted API, Anthropic, or Amazon Bedrock.

3. Reproduce the dataset comparison

The full comparison downloads GSM8K and an E5 encoder, embeds 4,000 training solutions, and queries an LLM four ways. Begin with the small smoke run:

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

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

When that works, use the full defaults:

python thinkretrieve/examples/compare_tts_vs_thinkretrieve.py \
  --backend openai --model qwen3:4b

See the 10-minute tutorial for Bedrock, Anthropic, hosted APIs, SciQ, procedural memory, troubleshooting, and the exact downloads performed by each example.

4. Try science reasoning + procedural memory + tool dispatch

The runnable science-agent case study uses a small local model to diagnose simulated cell-culture incidents, retrieve a private SOP mid-thought, and select an opaque tool route. On the frozen six-case screen, ThinkRetrieve reached 4/6 with Qwen 3.5 4B and 5/6 with 9B; plain, prompt-RAG, and more-thinking controls ranged from 0/6 to 2/6. The case study includes the exact runner, raw outputs, setup commands, and limitations.

Install the released library

pip install "thinkretrieve[faiss]"            # core + FAISS retrieval
pip install "thinkretrieve[faiss,anthropic]"  # add Anthropic
pip install "thinkretrieve[all]"              # all supported backends
BackendWhat you provide
Ollama / LM Studio / MLX / llama.cpp / vLLMOpenAI-compatible base URL and local model name
OpenAI / Together / Groq / DeepSeek / OpenRouterOpenAI-compatible base URL, model, and API key
AnthropicANTHROPIC_API_KEY and a Claude model
Amazon BedrockAWS credentials, region, and Bedrock model ID

Repository map

ThinkRetrieve/
├── README.md                       This repository guide
├── docs/                           GitHub Pages source
│   ├── index.html                  Project website
│   └── article/                    Hosted HTML + Markdown article + assets
└── thinkretrieve/                  Self-contained Python project
    ├── README.md                   Library and PyPI documentation
    ├── TUTORIAL.md                 Full guided walkthrough
    ├── LICENSE                     MIT license
    ├── pyproject.toml              Package metadata and dependencies
    ├── src/thinkretrieve/          Library source
    ├── examples/                   Offline, local, API, and dataset examples
    ├── benchmarks/science_agent/   Runnable agent case study + raw outputs
    ├── tests/                      Dependency-free core tests
    └── assets/                     Figures used by the documentation

Results

Across five reasoning models and four benchmarks in the paper, ThinkRetrieve outperforms standard thinking, sequential test-time scaling, static ICL, and random retrieval in every reported model–benchmark cell.

ThinkRetrieve results across models and benchmarks

Citation

@article{thinkretrieve2026,
  title   = {ThinkRetrieve: Retrieval-Augmented Reasoning Traces for Test-Time Scaling},
  author  = {Singh, Vaibhav and Ghosal, Soumya Suvra and Gharat, Sarvesh and
             Pal, Soumyabrata and Narayanam, Ramasuri and Manocha, Dinesh},
  journal = {arXiv preprint arXiv:2608.10928},
  year    = {2026},
  doi     = {10.48550/arXiv.2608.10928},
  url     = {https://arxiv.org/abs/2608.10928}
}

License

The library is released under the MIT License.