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

Uploaded Python 3

File details

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

File metadata

  • Download URL: cotlab-0.8.1.tar.gz
  • Upload date:
  • Size: 468.6 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.1.tar.gz
Algorithm Hash digest
SHA256 da8bf4ee981eb01075679ac08ec6d0db5ebfeceeca06e9d3690fba24fecacc56
MD5 e67445d51daccc221556218ca65eac96
BLAKE2b-256 77336bf4742a153f5b78f8b7d172bec539574b38b9ebe2c61be74a3ce40ba7be

See more details on using hashes here.

File details

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

File metadata

  • Download URL: cotlab-0.8.1-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.1-py3-none-any.whl
Algorithm Hash digest
SHA256 48b13da346a0acad7363b54c8e52bea9ff0bc4d5d7d94697d0297a52c7e1e712
MD5 2fbeda9ad249a7ab07a82a1e29e171a1
BLAKE2b-256 5e31ec8ec56428e426e63c6a13e128c50f26cbbfa3ef61ab137cebd6f8a6bc1f

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