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.3.tar.gz (468.7 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.3-py3-none-any.whl (199.1 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: cotlab-0.8.3.tar.gz
  • Upload date:
  • Size: 468.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.11.8 {"installer":{"name":"uv","version":"0.11.8","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.3.tar.gz
Algorithm Hash digest
SHA256 5621027d6ce4353a9cd1a0322e2c826ad60655e57d9d9daaf5731b3c6f423ec5
MD5 14357f6b0a8a1aa94f3d867885f7ae40
BLAKE2b-256 eebebc1dc3a2df7d83e50631eafb64b0d4a66cd3b1bd33d98ec38304ee461b51

See more details on using hashes here.

File details

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

File metadata

  • Download URL: cotlab-0.8.3-py3-none-any.whl
  • Upload date:
  • Size: 199.1 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.11.8 {"installer":{"name":"uv","version":"0.11.8","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.3-py3-none-any.whl
Algorithm Hash digest
SHA256 f27fcbcddebf03b0f1b9cddc6fad48d7f364e660df29f2e72940aa3003e3dfe3
MD5 acbd42e95dab17d3617e43d617c73a8c
BLAKE2b-256 a5af6fc2bf6350f8d2c5fc2be11e60700614c7d802d1c3dc41dae0d402efbad8

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