Skip to main content

CoTLab - Chain of Thought Research Toolkit for LLMs

Project description

CoTLab

Ask DeepWiki

A research toolkit for investigating Chain of Thought (CoT) reasoning, faithfulness, and mechanistic interpretability in Large Language Models.

Features

  • Experiments for CoT faithfulness, patching, logit‑lens, steering, and probing
  • Diverse prompt strategies (CoT, direct, adversarial, contrarian, few‑shot, etc.)
  • Configurable models, datasets, and backends (vLLM + Transformers)
  • Auto‑detect layers/heads at runtime
  • Hydra config for easy composition and multiruns


Installation

git clone https://github.com/huseyincavusbi/CoTLab.git
cd CoTLab
uv venv cotlab --python 3.11
source cotlab/bin/activate
uv pip install -e ".[dev]"

# GPU Setup:
# NVIDIA: uv pip install vllm
# AMD ROCm: ./scripts/cotlab-rocm.sh (uses Docker)

# AMD ROCm (Transformers backend): install ROCm PyTorch wheels
# uv pip install --reinstall --index-url https://download.pytorch.org/whl/rocm6.4 torch torchvision torchaudio

# Apple Silicon: requires Python 3.12 and vllm-metal plugin
# See docs/getting-started/installation.md for Metal setup instructions

See Installation Docs for detailed GPU setup.

Backend Compatibility

CoTLab supports two inference backends with different strengths:

1. vLLM Backend (High Performance)

Best for large-scale generation experiments.

  • Supported Experiments: cot_faithfulness, radiology
  • Supported Models: All text-only models (e.g., gemma_270m, medgemma_27b_text_it)
  • Limitation: Does NOT support activation patching or internal state access.
  • Note: Gemma 3 multimodal models (e.g., medgemma_4b_it) are currently incompatible with vLLM 0.12.0 due to architecture detection issues. Use transformers backend for these.

2. Transformers Backend (Full Access)

Best for mechanistic interpretability and activation patching.

  • Supported Experiments: ALL experiments.
  • Supported Models: ALL models.
  • Limitation: Slower.

To switch backends:

# Use vLLM (fast generation)
python -m cotlab.main backend=vllm ...

# Use Transformers (activation access)
python -m cotlab.main backend=transformers ...

Quick Start

# Run logit lens on MedGemma
python -m cotlab.main experiment=logit_lens model=medgemma_4b

# Find sycophancy heads
python -m cotlab.main experiment=sycophancy_heads model=medgemma_4b

# Test CoT ablation on pediatrics dataset
python -m cotlab.main experiment=cot_ablation dataset=pediatrics

# Compare prompt strategies
python -m cotlab.main -m prompt=chain_of_thought,direct_answer,sycophantic

Supported Models

CoTLab ships config files for some models, but in principle it supports any model that the selected backend can load. Mechanistic experiments can still fail for models with unusual architectures.

You can add a model config file for more control over hyperparameters, but you can also run any experiment by passing a Hugging Face model name directly.

# Use a built-in model config
python -m cotlab.main model=medgemma_4b

# Or pass any HF model name directly
python -m cotlab.main model.name=google/gemma-3-270m

If you prefer, you can pre-create model configs with cotlab-template, but this is optional because CoTLab can auto-generate model configs when running with model=org/repo-id.

Configuration

All configs auto-detect layers/heads at runtime. Override via CLI:

python -m cotlab.main \
    model=medgemma_4b \
    dataset=pediatrics \
    prompt=chain_of_thought \
    experiment.top_k=10

Acknowledgements

This research is conducted in collaboration with the Great Ormond Street Hospital DRIVE Unit.

Contributors

  • Huseyin Cavus — Core Contributor
  • Dr. Pavithra Rajendran — Machine Learning Lead, GOSH DRIVE
  • Sebin Sabu — Senior AI Scientist, GOSH DRIVE
  • Jaskaran Singh Kawatra — ML Engineer, GOSH DRIVE

License

MIT

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

cotlab-0.8.4.tar.gz (448.9 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

cotlab-0.8.4-py3-none-any.whl (199.2 kB view details)

Uploaded Python 3

File details

Details for the file cotlab-0.8.4.tar.gz.

File metadata

  • Download URL: cotlab-0.8.4.tar.gz
  • Upload date:
  • Size: 448.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.11.25 {"installer":{"name":"uv","version":"0.11.25","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for cotlab-0.8.4.tar.gz
Algorithm Hash digest
SHA256 11b5e1e984fafd456de1a1f9e60fcc450c7d3d9b477924eb284491a1d8c76685
MD5 5c1e1b2dae33e6e81c595ded382f0748
BLAKE2b-256 fb031c3a053c8ed1d3a9f82725cc32399f809228ebe2ec0106cef06850f3d144

See more details on using hashes here.

File details

Details for the file cotlab-0.8.4-py3-none-any.whl.

File metadata

  • Download URL: cotlab-0.8.4-py3-none-any.whl
  • Upload date:
  • Size: 199.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.11.25 {"installer":{"name":"uv","version":"0.11.25","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for cotlab-0.8.4-py3-none-any.whl
Algorithm Hash digest
SHA256 dffa679e6b345a92e28a324a335c5e636e5b0ad2811eb094ea158a8092dc1c0f
MD5 901caa144278ca84007f4972dcea6816
BLAKE2b-256 ac3a905822bcd013ca6de7f78e15d27c78a28a547e557daefdc749e68ca380e8

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page