chore: Update README for uv workflow

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,24 @@
+name: Bug Report
+description: File a bug report
+title: "[Bug]: "
+labels: ["bug"]
+body:
+  - type: textarea
+    id: description
+    attributes:
+      label: Bug Description
+      description: Describe what happened and what you expected
+      placeholder: What went wrong?
+    validations:
+      required: true
+
+  - type: textarea
+    id: steps
+    attributes:
+      label: Steps to reproduce
+      placeholder: |
+        1. 
+        2. 
+        3.
+    validations:
+      required: false

.github/ISSUE_TEMPLATE/feature_request.yml

@@ -0,0 +1,13 @@
+name: Feature Request
+description: Suggest an idea for this project
+title: "[Feature]: "
+labels: ["enhancement"]
+body:
+  - type: textarea
+    id: description
+    attributes:
+      label: Feature Description
+      description: What feature would you like to see?
+      placeholder: Describe the feature you'd like
+    validations:
+      required: true

.github/pull_request_template.md

@@ -0,0 +1,10 @@
+## What does this PR do?
+
+Describe your changes here.
+
+Fixes #
+
+## Checklist
+
+- [ ] I tested my changes
+- [ ] I reviewed my own code

.github/workflows/ci.yml

@@ -0,0 +1,76 @@
+name: Checks
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+permissions:
+  contents: read
+
+jobs:
+  lint:
+    name: Lint
+    runs-on: blacksmith-2vcpu-ubuntu-2404
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
+        with:
+          persist-credentials: false
+
+      - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
+        with:
+          python-version: "3.11"
+
+      - run: uv sync --locked --dev
+      - run: uv run ruff check .
+
+  format:
+    name: Format
+    runs-on: blacksmith-2vcpu-ubuntu-2404
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
+        with:
+          persist-credentials: false
+
+      - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
+        with:
+          python-version: "3.11"
+
+      - run: uv sync --locked --dev
+      - run: uv run ruff format --check .
+
+  test:
+    name: Test
+    runs-on: blacksmith-2vcpu-ubuntu-2404
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
+        with:
+          persist-credentials: false
+
+      - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
+        with:
+          python-version: "3.11"
+
+      - run: uv sync --locked --dev
+      - run: uv run pytest
+
+  build:
+    name: Build
+    runs-on: blacksmith-2vcpu-ubuntu-2404
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
+        with:
+          persist-credentials: false
+
+      - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
+        with:
+          python-version: "3.11"
+
+      - run: uv sync --locked --dev
+      - run: uv run python -m compileall activation_probe_mvp scripts tests
+      - run: uv build

.pre-commit-config.yaml

@@ -0,0 +1,20 @@
+repos:
+  - repo: local
+    hooks:
+      - id: lint
+        name: Lint
+        entry: uv run ruff check .
+        language: system
+        pass_filenames: false
+
+      - id: format
+        name: Format
+        entry: uv run ruff format --check .
+        language: system
+        pass_filenames: false
+
+      - id: build
+        name: Build
+        entry: uv build --out-dir /tmp/open-constitution-pre-commit-build
+        language: system
+        pass_filenames: false

README.md

@@ -12,19 +12,24 @@ This is a research MVP, not production-ready safety infrastructure.
 
 ## Install
 
+Install `uv` if you do not already have it:
+
 ```bash
-python3 -m venv .venv
-source .venv/bin/activate
+curl -LsSf https://astral.sh/uv/install.sh | sh
+```
 
-pip install -e .
+Then install the project dependencies:
+
+```bash
+uv sync
 ```
 
 ## Quick start
 
-The default model is intentionally small so the MVP runs locally:
+The default model is Gemma 4 E2B. It is smaller than many frontier models, but still needs substantial disk and memory:
 
 ```bash
-train-probe \
+uv run train-probe \
   --model_id google/gemma-4-E2B-it \
   --data_path data/examples.jsonl \
   --layer -4 \
@@ -34,7 +39,7 @@ train-probe \
 Then run guarded generation:
 
 ```bash
-guarded-generate \
+uv run guarded-generate \
   --model_id google/gemma-4-E2B-it \
   --probe_path ./probe_out/probe.pt \
   --config_path ./probe_out/config.json \
@@ -137,7 +142,7 @@ Disable chat templates with:
 Example Gemma 4 run:
 
 ```bash
-train-probe \
+uv run train-probe \
   --model_id google/gemma-4-E2B-it \
   --data_path data/examples.jsonl \
   --layer -4 \
@@ -147,7 +152,7 @@ train-probe \
 Then:
 
 ```bash
-guarded-generate \
+uv run guarded-generate \
   --model_id google/gemma-4-E2B-it \
   --probe_path ./probe_out_gemma4/probe.pt \
   --config_path ./probe_out_gemma4/config.json \
@@ -157,10 +162,10 @@ guarded-generate \
 Layer sweep:
 
 ```bash
-sweep-layers \
+uv run sweep-layers \
   --model_id google/gemma-4-E2B-it \
   --data_path data/examples.jsonl \
   --layers="-2,-4,-6,-8,-10,-12"
 ```
 
-Note: Gemma models may require accepting Google's license terms on Hugging Face before download.
+Note: Gemma models may require accepting Google's license terms on Hugging Face before download. `google/gemma-4-E2B-it` downloads about 10 GB of weights, so use a machine or cloud runtime with enough free disk space.

activation_probe_mvp/activations.py

@@ -79,10 +79,12 @@ def tokenize_text(tokenizer_or_processor, text: str, device: str, max_length: in
     }
 
     if max_length is not None:
-        kwargs.update({
-            "truncation": True,
-            "max_length": max_length,
-        })
+        kwargs.update(
+            {
+                "truncation": True,
+                "max_length": max_length,
+            }
+        )
 
     try:
         inputs = tokenizer_or_processor(text=text, **kwargs)
@@ -136,6 +138,6 @@ def extract_final_token_hidden_state(
     )
 
     selected = get_selected_hidden_state(outputs, layer)  # [batch, seq, hidden]
-    final_token = selected[:, -1, :]                     # [batch, hidden]
+    final_token = selected[:, -1, :]  # [batch, hidden]
 
     return final_token.float().detach().cpu()

activation_probe_mvp/data.py

@@ -3,7 +3,6 @@
 import json
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Iterable
 
 
 @dataclass

activation_probe_mvp/training.py

@@ -40,7 +40,7 @@ def train_linear_probe(
             total += float(loss.item())
 
         if epoch == 0 or (epoch + 1) % 20 == 0:
-            print(f"epoch={epoch+1:03d} loss={total / max(len(loader), 1):.4f}")
+            print(f"epoch={epoch + 1:03d} loss={total / max(len(loader), 1):.4f}")
 
     return probe
 

pyproject.toml

@@ -28,3 +28,15 @@ guarded-generate = "scripts.guarded_generate:main"
 
 [tool.setuptools.packages.find]
 include = ["activation_probe_mvp*", "scripts*"]
+
+[dependency-groups]
+dev = [
+    "pre-commit>=4.6.0",
+    "pytest>=8.0",
+    "ruff>=0.8",
+    "virtualenv<21",
+]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"

scripts/guarded_generate.py

@@ -81,7 +81,7 @@ def guarded_generate(
         risk = torch.sigmoid(probe(hidden)).item()
         smoothed = smoother.update(risk)
 
-        print(f"\n[probe] step={step+1} risk={risk:.3f} smoothed={smoothed:.3f}")
+        print(f"\n[probe] step={step + 1} risk={risk:.3f} smoothed={smoothed:.3f}")
 
         if smoothed >= pause_threshold:
             partial = "".join(generated_tokens)

tests/test_activations.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from activation_probe_mvp.activations import get_selected_hidden_state
+
+
+def test_get_selected_hidden_state_uses_hidden_states_first():
+    hidden_states = [torch.tensor([0]), torch.tensor([1])]
+    decoder_hidden_states = [torch.tensor([2])]
+    outputs = SimpleNamespace(
+        hidden_states=hidden_states,
+        decoder_hidden_states=decoder_hidden_states,
+    )
+
+    assert get_selected_hidden_state(outputs, -1) is hidden_states[-1]
+
+
+def test_get_selected_hidden_state_falls_back_to_decoder_hidden_states():
+    decoder_hidden_states = [torch.tensor([0]), torch.tensor([1])]
+    outputs = SimpleNamespace(hidden_states=None, decoder_hidden_states=decoder_hidden_states)
+
+    assert get_selected_hidden_state(outputs, 0) is decoder_hidden_states[0]
+
+
+def test_get_selected_hidden_state_errors_when_hidden_states_are_missing():
+    outputs = SimpleNamespace(hidden_states=None, decoder_hidden_states=None)
+
+    with pytest.raises(RuntimeError, match="did not include hidden_states"):
+        get_selected_hidden_state(outputs, 0)

tests/test_chat.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from activation_probe_mvp.chat import format_exchange, format_generation_prompt
+
+
+class TemplateOwner:
+    def __init__(self):
+        self.calls = []
+
+    def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False):
+        self.calls.append(
+            {
+                "messages": messages,
+                "tokenize": tokenize,
+                "add_generation_prompt": add_generation_prompt,
+            }
+        )
+        roles = ",".join(message["role"] for message in messages)
+        return f"template:{roles}:{add_generation_prompt}"
+
+
+class LegacyTemplateOwner:
+    def apply_chat_template(self, messages, tokenize=False):
+        roles = ",".join(message["role"] for message in messages)
+        return f"legacy:{roles}:{tokenize}"
+
+
+class ProcessorWrapper:
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+
+def test_format_exchange_falls_back_without_chat_template():
+    assert (
+        format_exchange(object(), prompt="Explain", response="Answer")
+        == "User:\nExplain\n\nAssistant:\nAnswer"
+    )
+
+
+def test_format_generation_prompt_falls_back_without_chat_template():
+    assert format_generation_prompt(object(), prompt="Explain") == "User:\nExplain\n\nAssistant:\n"
+
+
+def test_format_exchange_uses_direct_chat_template_owner():
+    owner = TemplateOwner()
+
+    assert (
+        format_exchange(owner, prompt="Explain", response="Answer")
+        == "template:user,assistant:False"
+    )
+    assert owner.calls == [
+        {
+            "messages": [
+                {"role": "user", "content": "Explain"},
+                {"role": "assistant", "content": "Answer"},
+            ],
+            "tokenize": False,
+            "add_generation_prompt": False,
+        }
+    ]
+
+
+def test_format_generation_prompt_uses_nested_tokenizer_template():
+    owner = TemplateOwner()
+    processor = ProcessorWrapper(owner)
+
+    assert format_generation_prompt(processor, prompt="Explain") == "template:user:True"
+    assert owner.calls[0]["messages"] == [{"role": "user", "content": "Explain"}]
+    assert owner.calls[0]["add_generation_prompt"] is True
+
+
+def test_format_exchange_supports_legacy_template_signature():
+    assert (
+        format_exchange(LegacyTemplateOwner(), prompt="Explain", response="Answer")
+        == "legacy:user,assistant:False"
+    )