cotlab 0.8.0__py3-none-any.whl

cotlab/main.py

@@ -0,0 +1,317 @@
+"""Main entry point for the CoT research framework."""
+
+from __future__ import annotations
+
+import inspect
+import os
+import random
+import re
+import sys
+import traceback
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Iterator, TextIO
+
+import hydra
+import numpy as np
+import torch
+from omegaconf import DictConfig, OmegaConf
+
+from .core import create_component
+from .experiment import ExperimentDocumenter
+from .logging import ExperimentLogger
+
+
+class _TeeStream:
+    """Mirror writes to multiple text streams."""
+
+    def __init__(self, *streams: TextIO):
+        self._streams = streams
+        self._primary = streams[0]
+
+    def write(self, data: str) -> int:
+        for stream in self._streams:
+            stream.write(data)
+        return len(data)
+
+    def flush(self) -> None:
+        for stream in self._streams:
+            stream.flush()
+
+    def isatty(self) -> bool:
+        isatty = getattr(self._primary, "isatty", None)
+        return bool(isatty()) if callable(isatty) else False
+
+    @property
+    def encoding(self) -> str:
+        return getattr(self._primary, "encoding", "utf-8")
+
+    def __getattr__(self, name: str):
+        # Delegate other stream attributes (e.g. fileno, buffer).
+        return getattr(self._primary, name)
+
+
+@contextmanager
+def _tee_stdout_stderr(log_path: Path) -> Iterator[None]:
+    """Tee stdout/stderr to the given log file."""
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+
+    original_stdout = sys.stdout
+    original_stderr = sys.stderr
+    log_file = open(log_path, "a", buffering=1, encoding="utf-8")
+
+    try:
+        sys.stdout = _TeeStream(original_stdout, log_file)
+        sys.stderr = _TeeStream(original_stderr, log_file)
+        yield
+    finally:
+        try:
+            sys.stdout.flush()
+            sys.stderr.flush()
+        finally:
+            sys.stdout = original_stdout
+            sys.stderr = original_stderr
+            log_file.close()
+
+
+def _safe_model_name(model_value: str) -> str:
+    return re.sub(r"[^a-zA-Z0-9]+", "_", model_value).strip("_").lower()
+
+
+def _repo_root() -> Path:
+    return Path(__file__).resolve().parents[2]
+
+
+def _normalize_hf_model_id(model_value: str) -> str:
+    """Normalize a HF model id: preserve org/, convert underscores to hyphens in the repo name.
+
+    google/medgemma_4b_it  ->  google/medgemma-4b-it
+    google/medgemma-4b-it  ->  google/medgemma-4b-it  (unchanged)
+    medgemma_4b_it         ->  medgemma_4b_it          (no slash, leave as-is)
+    """
+    if "/" in model_value:
+        org, repo = model_value.split("/", 1)
+        return f"{org}/{repo.replace('_', '-')}"
+    return model_value
+
+
+def _ensure_model_config_file(hf_id: str) -> None:
+    """Create conf/model/{org}/{repo}.yaml if it doesn't exist."""
+    hf_name = _normalize_hf_model_id(hf_id)
+    config_path = _repo_root() / "conf" / "model" / f"{hf_name}.yaml"
+    if config_path.exists():
+        return
+    config_path.parent.mkdir(parents=True, exist_ok=True)
+    content = "".join(
+        [
+            "# Auto-generated model config\n",
+            f"name: {hf_name}\n",
+            "max_new_tokens: 512\n",
+            "temperature: 0.7\n",
+            "top_p: 0.9\n",
+        ]
+    )
+    config_path.write_text(content)
+    print(f"[cotlab] Auto-generated model config: conf/model/{hf_name}.yaml")
+
+
+def _rewrite_hf_model_override(argv: list[str]) -> list[str]:
+    """Support `model=org/repo` by auto-generating conf/model/{org}/{repo}.yaml if missing."""
+    rewritten: list[str] = []
+
+    for arg in argv:
+        if arg.startswith("model="):
+            model_value = arg.split("=", 1)[1]
+            if "/" in model_value:
+                hf_id = _normalize_hf_model_id(model_value)
+                _ensure_model_config_file(hf_id)
+                rewritten.append(f"model={hf_id}")
+                continue
+        rewritten.append(arg)
+
+    return rewritten
+
+
+def _maybe_rewrite_argv() -> None:
+    """Rewrite argv in-place to support `model=<hf-id>` overrides and auto-config generation."""
+    if os.environ.get("COTLAB_DISABLE_HF_MODEL_REWRITE") == "1":
+        return
+    sys.argv[:] = _rewrite_hf_model_override(sys.argv)
+
+
+def set_seed(seed: int) -> None:
+    """Set random seeds for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def _extract_backend_load_kwargs(cfg_backend: DictConfig) -> dict:
+    backend_cfg = OmegaConf.to_container(cfg_backend, resolve=True)
+    keys = [
+        "load_in_4bit",
+        "load_in_8bit",
+        "bnb_4bit_quant_type",
+        "bnb_4bit_compute_dtype",
+        "bnb_4bit_use_double_quant",
+    ]
+    return {key: backend_cfg.get(key) for key in keys if backend_cfg.get(key) is not None}
+
+
+@hydra.main(version_base=None, config_path="../../conf", config_name="config")
+def _hydra_main(cfg: DictConfig) -> None:
+    """
+    Hydra-decorated experiment runner. Do not call directly — use main().
+
+    Uses Hydra for configuration management. Override configs via CLI:
+        python -m cotlab.main model.name=google/gemma-3-1b-it
+        python -m cotlab.main -m prompt=chain_of_thought,direct_answer  # multirun
+    """
+    # Get output directory from Hydra
+    output_dir = Path(hydra.core.hydra_config.HydraConfig.get().runtime.output_dir)
+    main_log_path = output_dir / "main.log"
+
+    with _tee_stdout_stderr(main_log_path):
+        try:
+            # Print config if verbose
+            if cfg.verbose:
+                print("=" * 60)
+                print("Configuration:")
+                print("=" * 60)
+                print(OmegaConf.to_yaml(cfg))
+                print("=" * 60)
+
+            # Set seed
+            set_seed(cfg.seed)
+
+            # Initialize logger
+            logger = ExperimentLogger(str(output_dir))
+            logger.log_config(cfg)
+
+            # Initialize experiment documenter and create initial EXPERIMENT.md
+            documenter = ExperimentDocumenter(cfg, output_dir)
+            doc_path = documenter.save()
+            print(f"Created experiment documentation: {doc_path}")
+
+            if cfg.dry_run:
+                print("Dry run - exiting without running experiment")
+                return
+
+            # Create components from config
+            print(f"Loading backend: {cfg.backend._target_}")
+            backend = create_component(cfg.backend)
+
+            print(f"Loading model: {cfg.model.name}")
+            backend.load_model(cfg.model.name, **_extract_backend_load_kwargs(cfg.backend))
+
+            print(f"Creating prompt strategy: {cfg.prompt.name}")
+            prompt_strategy = create_component(cfg.prompt)
+
+            print(f"Loading dataset: {cfg.dataset.name}")
+            dataset = create_component(cfg.dataset)
+
+            print(f"Creating experiment: {cfg.experiment.name}")
+            experiment = create_component(cfg.experiment)
+
+            # Run experiment
+            print("=" * 60)
+            print(f"Running experiment: {cfg.experiment.name}")
+            print("=" * 60)
+
+            try:
+                # Track experiment timing
+                import time
+
+                start_time = time.time()
+
+                # Generation kwargs (passed through to backend.generate_batch via experiments that accept **kwargs).
+                is_sampling = bool(cfg.model.temperature and cfg.model.temperature > 0)
+                gen_kwargs: dict = {
+                    "max_new_tokens": cfg.model.max_new_tokens,
+                    "temperature": cfg.model.temperature,
+                }
+                # top_p / top_k are only valid when sampling; suppress the HF warning otherwise.
+                if is_sampling:
+                    gen_kwargs["top_p"] = cfg.model.top_p
+
+                # HF Transformers uses `do_sample`; vLLM's SamplingParams does not accept it.
+                # For vLLM, greedy decoding is achieved by setting temperature=0.
+                if str(cfg.backend._target_).endswith("TransformersBackend"):
+                    gen_kwargs["do_sample"] = is_sampling
+
+                # num_samples is optional for some experiments; None means use all samples
+                num_samples = OmegaConf.select(cfg, "experiment.num_samples", default=None)
+                if isinstance(num_samples, int) and num_samples > 0:
+                    dataset_size = len(dataset)
+                    if num_samples > dataset_size:
+                        dataset_name = getattr(dataset, "name", type(dataset).__name__)
+                        print(
+                            "Warning: experiment.num_samples="
+                            f"{num_samples} exceeds dataset size ({dataset_size}) "
+                            f"for '{dataset_name}'. Using {dataset_size} samples."
+                        )
+                        num_samples = dataset_size
+
+                # Only pass kwargs that the experiment.run signature can accept.
+                # Some analysis experiments do not accept generation args.
+                run_sig = inspect.signature(experiment.run)
+                run_params = run_sig.parameters
+                accepts_var_kwargs = any(
+                    p.kind == inspect.Parameter.VAR_KEYWORD for p in run_params.values()
+                )
+
+                run_kwargs = {
+                    "backend": backend,
+                    "dataset": dataset,
+                    "prompt_strategy": prompt_strategy,
+                    "logger": logger,
+                }
+                extra_kwargs = dict(gen_kwargs)
+                extra_kwargs["num_samples"] = num_samples  # None = all samples, int = cap
+
+                if accepts_var_kwargs:
+                    run_kwargs.update(extra_kwargs)
+                else:
+                    run_kwargs.update({k: v for k, v in extra_kwargs.items() if k in run_params})
+
+                result = experiment.run(**run_kwargs)
+
+                # Calculate duration
+                duration = time.time() - start_time
+
+                # Save results
+                results_path = logger.save_results(result)
+                print(f"\nResults saved to: {results_path}")
+
+                # Print metrics summary
+                print("\nMetrics:")
+                for name, value in result.metrics.items():
+                    print(f"  {name}: {value}")
+
+                # Update EXPERIMENT.md with results
+                results_dict = {**result.metrics, "total_samples": len(result.raw_outputs)}
+                updated_doc = documenter.update_with_results(results_dict, duration)
+                documenter.save(updated_doc)
+                print(f"\nExperiment documentation updated: {doc_path}")
+
+            finally:
+                # Cleanup
+                backend.unload()
+                print("\nExperiment complete.")
+
+        except Exception:
+            # Ensure traceback is captured in main.log before Hydra handles it.
+            traceback.print_exc()
+            raise
+
+
+def main() -> None:
+    """Public entry point: rewrites argv for auto-config generation, then runs Hydra."""
+    _maybe_rewrite_argv()
+    _hydra_main()
+
+
+if __name__ == "__main__":
+    main()

cotlab/patching/__init__.py

@@ -0,0 +1,24 @@
+"""Activation patching module for causal interventions."""
+
+from .cache import ActivationCache
+from .hooks import HookManager
+from .interventions import (
+    Intervention,
+    InterventionType,
+    LayerImportance,
+    PatchingExperimentSpec,
+    ThoughtAnchor,
+)
+from .patcher import ActivationPatcher, PatchingResult
+
+__all__ = [
+    "HookManager",
+    "ActivationCache",
+    "ActivationPatcher",
+    "PatchingResult",
+    "InterventionType",
+    "Intervention",
+    "PatchingExperimentSpec",
+    "LayerImportance",
+    "ThoughtAnchor",
+]

cotlab/patching/cache.py

@@ -0,0 +1,141 @@
+"""Activation cache for storing and accessing layer outputs."""
+
+from dataclasses import dataclass, field
+from typing import Dict, Iterator, List, Optional, Tuple
+
+import torch
+
+
+@dataclass
+class ActivationCache:
+    """
+    Store activations from model forward passes.
+
+    This provides a clean interface for:
+    - Storing activations by layer index
+    - Retrieving full or sliced activations
+    - Managing GPU memory
+
+    Example:
+        >>> cache = ActivationCache()
+        >>> cache.store(0, layer_0_activations)
+        >>> cache.store(5, layer_5_activations)
+        >>> act = cache.get(0)  # Shape: [batch, seq_len, hidden_dim]
+        >>> sliced = cache.slice_tokens(0, (10, 20))  # Tokens 10-20
+    """
+
+    _cache: Dict[int, torch.Tensor] = field(default_factory=dict)
+    _metadata: Dict[str, any] = field(default_factory=dict)
+
+    def store(self, layer_idx: int, activation: torch.Tensor) -> None:
+        """
+        Store activation for a layer.
+
+        Args:
+            layer_idx: Layer index
+            activation: Activation tensor, typically [batch, seq_len, hidden_dim]
+        """
+        self._cache[layer_idx] = activation
+
+    def get(self, layer_idx: int) -> Optional[torch.Tensor]:
+        """
+        Retrieve activation for a layer.
+
+        Args:
+            layer_idx: Layer index
+
+        Returns:
+            Activation tensor or None if not cached
+        """
+        return self._cache.get(layer_idx)
+
+    def __getitem__(self, layer_idx: int) -> torch.Tensor:
+        """Allow dict-like access: cache[0]"""
+        if layer_idx not in self._cache:
+            raise KeyError(f"Layer {layer_idx} not in cache. Available: {self.layers}")
+        return self._cache[layer_idx]
+
+    def __contains__(self, layer_idx: int) -> bool:
+        """Check if layer is cached: 0 in cache"""
+        return layer_idx in self._cache
+
+    def __iter__(self) -> Iterator[int]:
+        """Iterate over cached layer indices."""
+        return iter(sorted(self._cache.keys()))
+
+    def __len__(self) -> int:
+        """Number of cached layers."""
+        return len(self._cache)
+
+    @property
+    def layers(self) -> List[int]:
+        """List of cached layer indices, sorted."""
+        return sorted(self._cache.keys())
+
+    @property
+    def shape(self) -> Optional[Tuple[int, ...]]:
+        """Shape of cached activations (assumed uniform)."""
+        if not self._cache:
+            return None
+        first = next(iter(self._cache.values()))
+        return tuple(first.shape)
+
+    def get_all(self) -> Dict[int, torch.Tensor]:
+        """Get all cached activations."""
+        return self._cache.copy()
+
+    def slice_tokens(self, layer_idx: int, token_range: Tuple[int, int]) -> torch.Tensor:
+        """
+        Get activations for specific token positions.
+
+        Args:
+            layer_idx: Layer index
+            token_range: (start, end) tuple, end is exclusive
+
+        Returns:
+            Sliced activation tensor
+        """
+        activation = self._cache[layer_idx]
+        start, end = token_range
+        return activation[:, start:end, :]
+
+    def slice_positions(self, layer_idx: int, positions: List[int]) -> torch.Tensor:
+        """
+        Get activations for specific non-contiguous positions.
+
+        Args:
+            layer_idx: Layer index
+            positions: List of token positions
+
+        Returns:
+            Activation tensor with shape [batch, len(positions), hidden_dim]
+        """
+        activation = self._cache[layer_idx]
+        return activation[:, positions, :]
+
+    def clear(self) -> None:
+        """Clear all cached activations to free memory."""
+        self._cache.clear()
+        torch.cuda.empty_cache()
+
+    def to_device(self, device: str) -> "ActivationCache":
+        """Move all cached activations to a device."""
+        new_cache = ActivationCache()
+        for layer_idx, activation in self._cache.items():
+            new_cache.store(layer_idx, activation.to(device))
+        return new_cache
+
+    def detach(self) -> "ActivationCache":
+        """Create a new cache with all activations detached."""
+        new_cache = ActivationCache()
+        for layer_idx, activation in self._cache.items():
+            new_cache.store(layer_idx, activation.detach().clone())
+        return new_cache
+
+    def set_metadata(self, key: str, value: any) -> None:
+        """Store metadata about the cache."""
+        self._metadata[key] = value
+
+    def get_metadata(self, key: str) -> any:
+        """Retrieve metadata."""
+        return self._metadata.get(key)