CoT Attribution Patching Prototype

Prototype implementation for generating an open-weight reasoning trace and analyzing sentence-to-sentence influence with attribution patching. The repo now supports both:

• local CLI experimentation
• a minimal GPU-hosted API for Runpod-style deployment

Scope

This repository currently implements:

• a Python library for generation plus attribution analysis
• a CLI prototype for end-to-end runs
• a FastAPI service for cloud GPU execution
• unit tests for sentence splitting, token alignment, summaries, mocked analysis, and API wiring

It does not yet include the full async web/queue stack from the larger product design.

Quick start

1. Install dependencies.
2. Run the CLI:

uv run python -m app.cli.run_prototype "Explain why the derivative of x^2 is 2x"

3. Inspect the emitted JSON in outputs/.

Runpod / Cloud GPU

For a dedicated Runpod pod, this repo now exposes a minimal API server:

uv run python -m app.cli.run_api

Endpoints:

• GET /healthz
• POST /api/warmup
• POST /api/analyze

Example request:

curl -X POST http://localhost:8000/api/analyze \
  -H 'content-type: application/json' \
  -d '{
    "question": "Explain why the derivative of x^2 is 2x",
    "max_new_tokens": 256,
    "max_trace_tokens": 512,
    "max_sentences": 24,
    "validate_top_k": 0
  }'

Recommended Runpod setup

• Start with a CUDA GPU pod, not serverless.
• Use the included Dockerfile.
• Keep deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B as the initial cloud model.
• Set DEVICE_PREFERENCE=cuda.
• Set PRELOAD_MODEL=true so the first request does not pay the full model load.
• Keep max_trace_tokens conservative at first; 512 is the safer initial setting.

Environment variables

• MODEL_NAME
• DEVICE_PREFERENCE
• MAX_TRACE_TOKENS
• MAX_SENTENCES
• TAKE_LOG
• PRELOAD_MODEL
• API_HOST
• API_PORT

Notes

• The analysis path requires attn_implementation="eager" and output_attentions=True.
• The default model is deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B.
• Local development on low-memory machines is suitable for tests and static checks, not for real attribution runs.
• Cloud GPU is the intended environment for real DeepSeek attribution analysis.