cotlab 0.8.0__py3-none-any.whl

cotlab/experiment/__init__.py

@@ -0,0 +1,315 @@
+"""Automatic experiment documentation generator.
+
+Generates EXPERIMENT.md files for every experiment run with:
+- Human-readable experiment description
+- Research context and questions
+- Full configuration
+- Reproduction commands
+- Results summary
+"""
+
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+
+class ExperimentDocumenter:
+    """Generates markdown documentation for experiments."""
+
+    def __init__(self, config: DictConfig, output_dir: Path):
+        self.config = config
+        self.output_dir = Path(output_dir)
+        self.start_time = datetime.now()
+
+    def _experiment_name(self) -> str:
+        return self.config.get("experiment", {}).get("name", "")
+
+    def _experiment_variants(self) -> list[dict]:
+        exp_cfg = self.config.get("experiment", {})
+        variants = exp_cfg.get("variants", [])
+        if isinstance(variants, (DictConfig, ListConfig)):
+            variants = OmegaConf.to_container(variants, resolve=True)
+        return list(variants) if variants else []
+
+    @staticmethod
+    def _format_variant_value(value: Any, base_value: str, default_label: str) -> str:
+        if isinstance(value, DictConfig):
+            value = OmegaConf.to_container(value, resolve=True)
+        if isinstance(value, dict):
+            value = value.get("name", default_label)
+        if value in ("base", "default"):
+            return base_value
+        return value if value else base_value
+
+    def generate_title(self) -> str:
+        """Generate human-readable experiment title from config."""
+        parts = []
+
+        exp_name = self._experiment_name()
+        if exp_name == "activation_compare":
+            variants = self._experiment_variants()
+            if variants:
+                run_names = [v.get("name", "run") for v in variants]
+                return f"Activation Compare: {' vs '.join(run_names)}"
+            return "Activation Compare"
+
+        # Dataset
+        dataset_name = self.config.get("dataset", {}).get("name", "unknown")
+        parts.append(dataset_name.capitalize())
+
+        # Reasoning mode
+        prompt_cfg = self.config.get("prompt", {})
+        if prompt_cfg.get("answer_first", False):
+            parts.append("Answer-First")
+        elif prompt_cfg.get("contrarian", False):
+            parts.append("Contrarian")
+        else:
+            parts.append("Standard")
+
+        # Few-shot
+        if not prompt_cfg.get("few_shot", True):
+            parts.append("Zero-Shot")
+
+        # Output format
+        output_fmt = prompt_cfg.get("output_format", "json")
+        if output_fmt != "json":
+            parts.append(f"({output_fmt.upper()})")
+
+        return " ".join(parts)
+
+    def infer_research_questions(self) -> list[str]:
+        """Infer research questions from configuration."""
+        questions = []
+        exp_name = self._experiment_name()
+
+        if exp_name == "activation_compare":
+            return [
+                "How do residual stream activations differ across runs and datasets?",
+                "Which layers show the largest activation divergence between runs?",
+                "Do activation differences align with task or prompt changes?",
+            ]
+
+        prompt_cfg = self.config.get("prompt", {})
+
+        # Few-shot ablation
+        if not prompt_cfg.get("few_shot", True):
+            questions.append("Does the model perform well without few-shot examples (zero-shot)?")
+
+        # Reasoning mode
+        if prompt_cfg.get("contrarian", False):
+            questions.append("Does skeptical/contrarian reasoning improve diagnostic accuracy?")
+        elif prompt_cfg.get("answer_first", False):
+            questions.append(
+                'Does "answer first, then justify" reasoning order affect performance?'
+            )
+
+        # Output format
+        output_fmt = prompt_cfg.get("output_format", "json")
+        if output_fmt != "json":
+            questions.append(
+                f"How does {output_fmt.upper()} output format affect parsing and accuracy?"
+            )
+
+        # Default question if none inferred
+        if not questions:
+            questions.append("Standard baseline experiment for comparison.")
+
+        return questions
+
+    def generate_reproduction_command(self) -> str:
+        """Generate exact CLI command to reproduce experiment."""
+        parts = ["python -m cotlab.main"]
+
+        # Add overrides from config
+        exp_cfg = self.config.get("experiment", {})
+        prompt_cfg = self.config.get("prompt", {})
+        dataset_cfg = self.config.get("dataset", {})
+
+        exp_name = exp_cfg.get("name")
+        if exp_name:
+            parts.append(f"experiment={exp_name}")
+        if exp_cfg.get("num_samples") is not None:
+            parts.append(f"experiment.num_samples={exp_cfg['num_samples']}")
+        if exp_cfg.get("seed") is not None:
+            parts.append(f"experiment.seed={exp_cfg['seed']}")
+
+        # Prompt selection
+        if "name" in prompt_cfg:
+            parts.append(f"prompt={prompt_cfg['name']}")
+
+        # Prompt parameters
+        if prompt_cfg.get("contrarian", False):
+            parts.append("prompt.contrarian=true")
+        if prompt_cfg.get("answer_first", False):
+            parts.append("prompt.answer_first=true")
+        if not prompt_cfg.get("few_shot", True):
+            parts.append("prompt.few_shot=false")
+
+        output_fmt = prompt_cfg.get("output_format", "json")
+        if output_fmt != "json":
+            parts.append(f"prompt.output_format={output_fmt}")
+
+        # Dataset
+        if "name" in dataset_cfg:
+            parts.append(f"dataset={dataset_cfg['name']}")
+
+        return " \\\n  ".join(parts)
+
+    def create_initial_doc(self) -> str:
+        """Create initial experiment documentation (before results)."""
+        title = self.generate_title()
+        questions = self.infer_research_questions()
+        repro_cmd = self.generate_reproduction_command()
+
+        # Configuration summary
+        prompt_cfg = self.config.get("prompt", {})
+        dataset_cfg = self.config.get("dataset", {})
+
+        reasoning_mode = "Standard"
+        if prompt_cfg.get("answer_first", False):
+            reasoning_mode = "Answer-First"
+        elif prompt_cfg.get("contrarian", False):
+            reasoning_mode = "Contrarian (skeptical)"
+
+        few_shot = "Yes" if prompt_cfg.get("few_shot", True) else "No (zero-shot)"
+        output_fmt = prompt_cfg.get("output_format", "json").upper()
+        dataset_name = dataset_cfg.get("name", "unknown")
+
+        doc = f"""# Experiment: {title}
+
+**Status:** Running
+**Started:** {self.start_time.strftime("%Y-%m-%d %H:%M:%S")}
+
+## Research Questions
+
+"""
+        for i, question in enumerate(questions, 1):
+            doc += f"{i}. {question}\n"
+
+        doc += f"""
+## Configuration
+
+**Prompt Strategy:** {prompt_cfg.get("name", "unknown").capitalize()}
+**Reasoning Mode:** {reasoning_mode}
+**Few-Shot Examples:** {few_shot}
+**Output Format:** {output_fmt}
+**Dataset:** {dataset_name}
+"""
+
+        exp_name = self._experiment_name()
+        if exp_name == "activation_compare":
+            variants = self._experiment_variants()
+            if variants:
+                doc += "\n**Variants:**\n"
+                for variant in variants:
+                    variant_name = variant.get("name", "run")
+                    variant_dataset = variant.get("dataset", "base")
+                    variant_prompt = variant.get("prompt", "base")
+                    variant_samples = variant.get("num_samples", "default")
+                    variant_seed = variant.get("seed", "default")
+                    variant_dataset = self._format_variant_value(
+                        variant_dataset, dataset_name, "dataset"
+                    )
+                    variant_prompt = self._format_variant_value(
+                        variant_prompt, prompt_cfg.get("name", "prompt"), "prompt"
+                    )
+                    doc += (
+                        f"- {variant_name}: dataset={variant_dataset}, "
+                        f"prompt={variant_prompt}, samples={variant_samples}, seed={variant_seed}\n"
+                    )
+
+        doc += f"""
+<details>
+<summary>Full Configuration (YAML)</summary>
+
+```yaml
+{OmegaConf.to_yaml(self.config)}
+```
+</details>
+
+## Reproduce
+
+```bash
+{repro_cmd}
+```
+
+## Results
+
+_Results will be added after experiment completes..._
+"""
+
+        return doc
+
+    def update_with_results(
+        self, results: Optional[Dict[str, Any]] = None, duration_seconds: Optional[float] = None
+    ) -> str:
+        """Update documentation with results after completion."""
+        # Read existing doc
+        doc_path = self.output_dir / "EXPERIMENT.md"
+        if doc_path.exists():
+            doc = doc_path.read_text()
+        else:
+            doc = self.create_initial_doc()
+
+        # Update status
+        doc = doc.replace("**Status:** Running", "**Status:** Completed")
+
+        # Add duration if provided
+        if duration_seconds is not None:
+            minutes = int(duration_seconds // 60)
+            seconds = int(duration_seconds % 60)
+            duration_str = (
+                f"{minutes} minutes {seconds} seconds" if minutes > 0 else f"{seconds} seconds"
+            )
+            doc = doc.replace(
+                f"**Started:** {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}",
+                f"**Started:** {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}  \n**Duration:** {duration_str}",
+            )
+
+        # Update results section
+        if results:
+            results_md = self._format_results(results)
+            doc = doc.replace(
+                "## Results\n\n_Results will be added after experiment completes..._",
+                f"## Results\n\n{results_md}",
+            )
+
+        return doc
+
+    def _format_results(self, results: Dict[str, Any]) -> str:
+        """Format results dictionary as markdown."""
+        md = ""
+        is_activation_compare = "num_runs" in results and "pair_count" in results
+
+        # Basic metrics
+        if "accuracy" in results:
+            md += f"- **Accuracy:** {results['accuracy']:.1%}\n"
+        if "total_samples" in results and not is_activation_compare:
+            md += f"- **Samples Processed:** {results['total_samples']}\n"
+        if "parse_failures" in results:
+            md += f"- **Parse Failures:** {results['parse_failures']}\n"
+        if "avg_time" in results:
+            md += f"- **Average Time per Sample:** {results['avg_time']:.2f}s\n"
+
+        # Additional metrics
+        for key, value in results.items():
+            if key not in ["accuracy", "total_samples", "parse_failures", "avg_time"]:
+                if isinstance(value, float):
+                    md += f"- **{key.replace('_', ' ').title()}:** {value:.3f}\n"
+                else:
+                    md += f"- **{key.replace('_', ' ').title()}:** {value}\n"
+
+        return md
+
+    def save(self, content: Optional[str] = None) -> Path:
+        """Save experiment documentation to file."""
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        doc_path = self.output_dir / "EXPERIMENT.md"
+
+        if content is None:
+            content = self.create_initial_doc()
+
+        doc_path.write_text(content)
+        return doc_path

cotlab/experiments/__init__.py

@@ -0,0 +1,43 @@
+"""Experiments module."""
+
+from .activation_compare import ActivationCompareExperiment
+from .activation_patching import ActivationPatchingExperiment
+from .attention_analysis import AttentionAnalysisExperiment
+from .classification import ClassificationExperiment
+from .composite_shift_detector import CompositeShiftDetectorExperiment
+from .cot_ablation import CoTAblationExperiment
+from .cot_faithfulness import CoTFaithfulnessExperiment
+from .cot_heads import CoTHeadsExperiment
+from .full_layer_cot import FullLayerCoTExperiment
+from .full_layer_patching import FullLayerPatchingExperiment
+from .h_neuron_analysis import HNeuronAnalysisExperiment
+from .logit_lens import LogitLensExperiment
+from .multi_head_cot import MultiHeadCoTExperiment
+from .multi_head_patching import MultiHeadPatchingExperiment
+from .probing_classifier import ProbingClassifierExperiment
+from .residual_norm_ood import ResidualNormOODExperiment
+from .sae_feature_analysis import SAEFeatureAnalysisExperiment
+from .steering_vectors import SteeringVectorsExperiment
+from .sycophancy_heads import SycophancyHeadsExperiment
+
+__all__ = [
+    "HNeuronAnalysisExperiment",
+    "CompositeShiftDetectorExperiment",
+    "CoTAblationExperiment",
+    "CoTFaithfulnessExperiment",
+    "ActivationPatchingExperiment",
+    "ActivationCompareExperiment",
+    "AttentionAnalysisExperiment",
+    "ProbingClassifierExperiment",
+    "ClassificationExperiment",
+    "ResidualNormOODExperiment",
+    "SAEFeatureAnalysisExperiment",
+    "SycophancyHeadsExperiment",
+    "MultiHeadPatchingExperiment",
+    "FullLayerPatchingExperiment",
+    "SteeringVectorsExperiment",
+    "CoTHeadsExperiment",
+    "LogitLensExperiment",
+    "MultiHeadCoTExperiment",
+    "FullLayerCoTExperiment",
+]

cotlab/experiments/activation_compare.py

@@ -0,0 +1,290 @@
+"""Activation Compare experiment — collect mean residual-stream vectors.
+
+One run = one condition (dataset + prompt settings).
+Saves per-layer mean activation vectors to results.json so that two or more
+runs can be compared offline with ``scripts/compare_activations.py``.
+
+Design follows logit_lens.py: hooks project inside the callback and move
+tensors to CPU immediately to avoid GPU page-faults on long reports.
+"""
+
+from typing import Any, Dict, List, Optional
+
+import torch
+from tqdm import tqdm
+
+from ..backends.base import InferenceBackend
+from ..core.base import BaseExperiment, ExperimentResult
+from ..core.registry import Registry
+from ..datasets.loaders import BaseDataset
+from ..logging import ExperimentLogger
+
+
+@Registry.register_experiment("activation_compare")
+class ActivationCompareExperiment(BaseExperiment):
+    """
+    Collect layer-wise mean residual-stream activations for one condition.
+
+    Forward-passes N samples through the model.  At each layer a lightweight
+    hook captures the hidden state at the last token (or mean-pooled across
+    all tokens) and moves it to CPU immediately.  After all samples the
+    per-layer mean vector is computed and saved to results.json.
+
+    Two saved runs can then be compared with ``scripts/compare_activations.py``
+    which computes cosine-similarity and L2-distance profiles per layer.
+    """
+
+    def __init__(
+        self,
+        name: str = "activation_compare",
+        description: str = "Collect mean layer activations for representational comparison",
+        layer_stride: int = 2,
+        num_samples: Optional[int] = None,
+        pooling: str = "last_token",  # "last_token" | "mean"
+        max_input_tokens: int = 1024,
+        seed: int = 42,
+        answer_cue: str = "\n\nAnswer:",  # appended so last position mirrors logit_lens
+        batch_size: int = 1,
+        # Legacy fields kept so old YAML configs don't break
+        layers: Optional[List[int]] = None,
+        variants: Optional[List[Dict[str, Any]]] = None,
+        comparison_mode: str = "pairwise",
+        store_per_layer: bool = True,
+        log_samples: bool = False,
+        **kwargs,
+    ):
+        self._name = name
+        self.description = description
+        self.layer_stride = layer_stride
+        self.num_samples = num_samples
+        self.pooling = pooling
+        self.max_input_tokens = max_input_tokens
+        self.seed = seed
+        self.answer_cue = answer_cue
+        self.batch_size = max(1, int(batch_size))
+        # Legacy fields silently ignored — kept for backward compat
+        self._layers_legacy = layers
+        self._variants_legacy = variants
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _resolve_layers(self, backend: InferenceBackend) -> List[int]:
+        all_layers = list(range(backend.hook_manager.num_layers))
+        return all_layers[:: self.layer_stride]
+
+    def _pool(self, tensor: torch.Tensor) -> torch.Tensor:
+        """tensor: [seq_len, hidden_size] → [hidden_size]"""
+        if self.pooling == "last_token":
+            return tensor[-1]
+        elif self.pooling == "mean":
+            return tensor.mean(dim=0)
+        else:
+            raise ValueError(f"Unknown pooling: {self.pooling}")
+
+    def _pool_batch(self, tensor: torch.Tensor) -> torch.Tensor:
+        """tensor: [B, seq_len, hidden_size] → [B, hidden_size]"""
+        if self.pooling == "last_token":
+            return tensor[:, -1, :]  # left-padded: last position = last real token
+        elif self.pooling == "mean":
+            return tensor.mean(dim=1)
+        else:
+            raise ValueError(f"Unknown pooling: {self.pooling}")
+
+    # ------------------------------------------------------------------
+    # Main entry point
+    # ------------------------------------------------------------------
+
+    def run(
+        self,
+        backend: InferenceBackend,
+        dataset: BaseDataset,
+        prompt_strategy: Any,
+        logger: Optional[ExperimentLogger] = None,
+        **kwargs,
+    ) -> ExperimentResult:
+        """Collect mean residual-stream activations for one dataset condition."""
+
+        target_layers = self._resolve_layers(backend)
+        tokenizer = backend._tokenizer
+
+        print(f"Model : {backend.model_name}")
+        print(f"Layers ({len(target_layers)}): {target_layers}")
+        print(f"Pooling: {self.pooling}")
+        print(f"Batch size: {self.batch_size}")
+
+        if self.num_samples is None:
+            samples = list(dataset)
+        else:
+            samples = dataset.sample(self.num_samples, seed=self.seed)
+        n = len(samples)
+        print(f"Samples: {n}\n")
+
+        # Accumulators: layer_idx → running sum tensor (float32, CPU)
+        layer_sums: Dict[int, torch.Tensor] = {}
+        layer_sq_sums: Dict[int, torch.Tensor] = {}
+        layer_counts: Dict[int, int] = {}
+        processed = 0
+
+        # Chunk into batches
+        batches = [
+            samples[i : i + self.batch_size] for i in range(0, len(samples), self.batch_size)
+        ]
+
+        for batch in tqdm(batches, desc="Activation collect"):
+            prompt_strs = []
+            for sample in batch:
+                prompt_input = {
+                    "text": sample.text,
+                    "question": sample.text,
+                    "report": sample.text,
+                    "metadata": sample.metadata or {},
+                }
+                prompt_strs.append(prompt_strategy.build_prompt(prompt_input) + self.answer_cue)
+
+            B = len(prompt_strs)
+
+            if B == 1:
+                # Single-sample path — no padding overhead
+                tokens = tokenizer(
+                    prompt_strs[0],
+                    return_tensors="pt",
+                    truncation=True,
+                    max_length=self.max_input_tokens,
+                ).to(backend.device)
+            else:
+                # Batched path — left-pad so position [-1] = last real token
+                orig_side = tokenizer.padding_side
+                tokenizer.padding_side = "left"
+                if tokenizer.pad_token_id is None:
+                    tokenizer.pad_token_id = tokenizer.eos_token_id
+                tokens = tokenizer(
+                    prompt_strs,
+                    return_tensors="pt",
+                    truncation=True,
+                    max_length=self.max_input_tokens,
+                    padding=True,
+                ).to(backend.device)
+                tokenizer.padding_side = orig_side
+
+                # Compute position_ids that skip padding so positional embeddings
+                # are identical to the non-padded single-sample case.
+                # For left-padded input: e.g. [PAD, PAD, t0, t1, t2] → positions [0,0,0,1,2]
+                attention_mask = tokens["attention_mask"]
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids = position_ids.masked_fill(attention_mask == 0, 0)
+                tokens["position_ids"] = position_ids
+
+            # Capture hidden states in hooks; move to CPU immediately
+            layer_vecs: Dict[int, torch.Tensor] = {}  # layer_idx -> [B, hidden] CPU float32
+
+            def make_hook(layer_idx: int):
+                def hook(module, inp, output):
+                    tensor = output[0] if isinstance(output, tuple) else output
+                    with torch.no_grad():
+                        if B == 1:
+                            # tensor: [1, seq_len, hidden]
+                            vec = self._pool(tensor[0]).unsqueeze(0).cpu().float()  # [1, hidden]
+                        else:
+                            # tensor: [B, seq_len, hidden]
+                            vec = self._pool_batch(tensor).cpu().float()  # [B, hidden]
+                    layer_vecs[layer_idx] = vec
+                    return output
+
+                return hook
+
+            handles = []
+            for layer_idx in target_layers:
+                if layer_idx < backend.hook_manager.num_layers:
+                    mod = backend.hook_manager.get_residual_module(layer_idx)
+                    handles.append(mod.register_forward_hook(make_hook(layer_idx)))
+
+            try:
+                with torch.no_grad():
+                    backend._model(**tokens)
+            except Exception as e:
+                tqdm.write(f"  [skip] batch starting at {batch[0].idx}: {type(e).__name__}: {e}")
+                torch.cuda.empty_cache()
+                for h in handles:
+                    h.remove()
+                continue
+            finally:
+                for h in handles:
+                    h.remove()
+
+            # Accumulate — iterate over each sample in the batch
+            for layer_idx, vecs in layer_vecs.items():  # vecs: [B, hidden]
+                for b in range(B):
+                    vec = vecs[b]  # [hidden]
+                    if layer_idx not in layer_sums:
+                        layer_sums[layer_idx] = torch.zeros_like(vec)
+                        layer_sq_sums[layer_idx] = torch.zeros_like(vec)
+                        layer_counts[layer_idx] = 0
+                    layer_sums[layer_idx] += vec
+                    layer_sq_sums[layer_idx] += vec**2
+                    layer_counts[layer_idx] += 1
+
+            del layer_vecs
+            torch.cuda.empty_cache()
+            processed += B
+
+        # --- Compute statistics -----------------------------------------
+        mean_activations: Dict[int, List[float]] = {}
+        activation_norm: Dict[int, float] = {}
+        activation_std: Dict[int, float] = {}
+
+        for layer_idx in target_layers:
+            cnt = layer_counts.get(layer_idx, 0)
+            if cnt == 0:
+                continue
+            mean_vec = layer_sums[layer_idx] / cnt  # [hidden]
+            var_vec = layer_sq_sums[layer_idx] / cnt - mean_vec**2
+            std_val = float(var_vec.clamp(min=0).sqrt().mean().item())
+            norm_val = float(mean_vec.norm().item())
+            mean_activations[layer_idx] = mean_vec.tolist()
+            activation_norm[layer_idx] = round(norm_val, 4)
+            activation_std[layer_idx] = round(std_val, 6)
+
+        # --- Print summary -----------------------------------------------
+        print("\n" + "=" * 60)
+        print("ACTIVATION COLLECT SUMMARY")
+        print("=" * 60)
+        print(f"Processed samples : {processed} / {n}")
+        print(f"{'Layer':>6}  {'Norm':>10}  {'Std':>10}")
+        print("-" * 32)
+        for layer_idx in target_layers:
+            if layer_idx in activation_norm:
+                print(
+                    f"{layer_idx:>6}  {activation_norm[layer_idx]:>10.2f}  {activation_std[layer_idx]:>10.6f}"
+                )
+        print("=" * 60)
+
+        return ExperimentResult(
+            experiment_name=self.name,
+            model_name=backend.model_name,
+            prompt_strategy=(
+                prompt_strategy.name if hasattr(prompt_strategy, "name") else "custom"
+            ),
+            metrics={
+                "num_samples": processed,
+                "pooling": self.pooling,
+                "layer_stride": self.layer_stride,
+                "activation_norms": activation_norm,
+                "activation_std": activation_std,
+            },
+            raw_outputs={
+                "mean_activations_per_layer": mean_activations,
+            },
+            metadata={
+                "target_layers": target_layers,
+                "pooling": self.pooling,
+                "num_samples": processed,
+                "seed": self.seed,
+                "answer_cue": self.answer_cue,
+            },
+        )