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.2.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.2-py3-none-any.whl (199.1 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: cotlab-0.8.2.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.2.tar.gz
Algorithm Hash digest
SHA256 f36b397cdb04697927f3cea43da270928b12be0557dfaf1c5297cdebce7619ed
MD5 ea27153317a7f14987b16608daa31161
BLAKE2b-256 0d202718d5fe8cf99759bb8afab5d19530cb69671abe55e53a88ef0e66fe3daa

See more details on using hashes here.

File details

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

File metadata

  • Download URL: cotlab-0.8.2-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.2-py3-none-any.whl
Algorithm Hash digest
SHA256 9aa5ea69db3f09becd486f7fe444b8e539980f24c3a866fc6dd635a3abf41458
MD5 83b69e489c354e4273f12d4a151528e5
BLAKE2b-256 9e4078ecbdb522ad904fe7f46374d3180abbce09fb6ad02e1d837246f0df976e

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