wisent 0.1.1__py3-none-any.whl → 0.5.1__py3-none-any.whl

wisent/core/user_model_config.py

@@ -0,0 +1,158 @@
+"""
+User-defined model configuration storage and retrieval.
+Handles models that aren't explicitly supported by storing user-provided configurations.
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Dict, Optional, Any
+from enum import Enum
+
+
+class ModelArchitecture(Enum):
+    """Supported model architectures for layer access."""
+    LLAMA_STYLE = "llama_style"  # model.layers.{idx}
+    GPT2_STYLE = "gpt2_style"    # transformer.h.{idx}
+    MPT_STYLE = "mpt_style"      # transformer.blocks.{idx}
+    CUSTOM = "custom"             # User provides full path template
+
+
+class UserModelConfig:
+    """Manages user-defined model configurations."""
+    
+    def __init__(self):
+        # Store config in user's home directory
+        self.config_dir = Path.home() / ".wisent-guard"
+        self.config_file = self.config_dir / "user_model_configs.json"
+        self.configs = self._load_configs()
+    
+    def _load_configs(self) -> Dict[str, Any]:
+        """Load existing configurations from file."""
+        if self.config_file.exists():
+            try:
+                with open(self.config_file, 'r') as f:
+                    return json.load(f)
+            except Exception:
+                return {}
+        return {}
+    
+    def _save_configs(self) -> None:
+        """Save configurations to file."""
+        self.config_dir.mkdir(exist_ok=True)
+        with open(self.config_file, 'w') as f:
+            json.dump(self.configs, f, indent=2)
+    
+    def has_config(self, model_name: str) -> bool:
+        """Check if we have a configuration for this model."""
+        return model_name in self.configs
+    
+    def get_config(self, model_name: str) -> Optional[Dict[str, Any]]:
+        """Get configuration for a model."""
+        return self.configs.get(model_name)
+    
+    def save_config(self, model_name: str, config: Dict[str, Any]) -> None:
+        """Save configuration for a model."""
+        self.configs[model_name] = config
+        self._save_configs()
+    
+    def get_prompt_tokens(self, model_name: str) -> Optional[Dict[str, str]]:
+        """Get user and assistant tokens for a model."""
+        config = self.get_config(model_name)
+        if config:
+            return {
+                "user_token": config.get("user_token"),
+                "assistant_token": config.get("assistant_token"),
+                "system_token": config.get("system_token"),  # Optional
+                "format_template": config.get("format_template")  # Optional custom template
+            }
+        return None
+    
+    def get_layer_access_info(self, model_name: str) -> Optional[Dict[str, Any]]:
+        """Get layer access information for a model."""
+        config = self.get_config(model_name)
+        if config:
+            return {
+                "architecture": config.get("architecture"),
+                "layer_path_template": config.get("layer_path_template"),
+                "custom_layer_accessor": config.get("custom_layer_accessor")
+            }
+        return None
+    
+    def prompt_and_save_config(self, model_name: str) -> Dict[str, Any]:
+        """
+        Interactively prompt user for model configuration.
+        This should be called from the CLI when an unknown model is encountered.
+        """
+        print(f"\n⚠️  Model '{model_name}' is not recognized.")
+        print("We need some information to properly support this model.\n")
+        
+        config = {"model_name": model_name}
+        
+        # Prompt for tokens
+        print("1. Chat Format Tokens")
+        print("   These are the special tokens your model uses to distinguish user and assistant messages.")
+        print("   Examples:")
+        print("   - Llama 3: <|start_header_id|>user<|end_header_id|> and <|start_header_id|>assistant<|end_header_id|>")
+        print("   - ChatGPT: <|im_start|>user and <|im_start|>assistant")
+        print("   - Alpaca: ### Human: and ### Assistant:")
+        
+        config["user_token"] = input("\n   Enter the user token/prefix: ").strip()
+        config["assistant_token"] = input("   Enter the assistant token/prefix: ").strip()
+        
+        # Optional system token
+        system_token = input("   Enter the system token/prefix (press Enter to skip): ").strip()
+        if system_token:
+            config["system_token"] = system_token
+        
+        # Model architecture for layer access
+        print("\n2. Model Architecture")
+        print("   How are the transformer layers accessed in this model?")
+        print("   1. Llama-style: model.layers.{idx}")
+        print("   2. GPT2-style: transformer.h.{idx}")
+        print("   3. MPT-style: transformer.blocks.{idx}")
+        print("   4. Custom (you'll provide the template)")
+        
+        while True:
+            choice = input("\n   Select architecture (1-4): ").strip()
+            if choice == "1":
+                config["architecture"] = ModelArchitecture.LLAMA_STYLE.value
+                config["layer_path_template"] = "model.layers.{idx}"
+                break
+            elif choice == "2":
+                config["architecture"] = ModelArchitecture.GPT2_STYLE.value
+                config["layer_path_template"] = "transformer.h.{idx}"
+                break
+            elif choice == "3":
+                config["architecture"] = ModelArchitecture.MPT_STYLE.value
+                config["layer_path_template"] = "transformer.blocks.{idx}"
+                break
+            elif choice == "4":
+                config["architecture"] = ModelArchitecture.CUSTOM.value
+                template = input("   Enter the layer path template (use {idx} for layer index): ").strip()
+                config["layer_path_template"] = template
+                break
+            else:
+                print("   Invalid choice. Please enter 1, 2, 3, or 4.")
+        
+        # Optional: custom format template
+        print("\n3. Custom Format Template (Optional)")
+        print("   If your model requires a specific prompt format beyond simple token prefixes,")
+        print("   you can provide a template. Use {user_message} and {assistant_message} as placeholders.")
+        print("   Example: '<|system|>\\nYou are a helpful assistant\\n{user_message}\\n{assistant_message}'")
+        
+        custom_template = input("\n   Enter custom template (press Enter to skip): ").strip()
+        if custom_template:
+            config["format_template"] = custom_template
+        
+        # Save the configuration
+        self.save_config(model_name, config)
+        
+        print(f"\n✅ Configuration saved for {model_name}")
+        print(f"   Config location: {self.config_file}")
+        
+        return config
+
+
+# Global instance for easy access
+user_model_configs = UserModelConfig()

wisent/opti/core/atoms.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import optuna
+
+__all__ = [
+    "Direction",
+    "HPOConfig",
+    "HPORun",
+    "BaseOptimizer",
+]
+
+Direction = Literal["maximize", "minimize"]
+
+
+@dataclass(slots=True, frozen=True)
+class HPOConfig:
+    """
+    Configuration for hyperparameter optimization (HPO) using Optuna.
+
+    attributes:
+        n_trials:
+            number of trials (ignored if timeout is reached).
+        direction:
+            global default direction ("maximize" or "minimize").
+        sampler:
+            one of {"tpe", "random", "cmaes"} or None for Optuna default.
+        pruner:
+            one of {"nop", "median", "sha", "asha", "hyperband"} or None for default.
+        timeout:
+            optional global seconds budget.
+        study_name:
+            optional persistent study name.
+        storage:
+            Optuna storage URL (e.g., sqlite:///file.db) for persistence.
+        seed:
+            sampler seed for reproducibility.
+        load_if_exists:
+            reuse persisted study if it already exists (when storage+study_name set).
+    """
+    n_trials: int = 100
+    direction: Direction = "maximize"
+    sampler: str | None = "tpe"
+    pruner: str | None = "asha"
+    timeout: int | None = None
+    storage: str | None = None
+    study_name: str | None = None
+    seed: int | None = 42
+    load_if_exists: bool = True
+
+
+@dataclass(slots=True, frozen=True)
+class HPORun:
+    """
+    Result of an HPO run.
+    """
+    study: optuna.Study
+    best_params: dict[str, Any]
+    best_value: float
+
+
+class BaseOptimizer(ABC):
+    """
+    Base class for building task-agnostic Optuna optimizers.
+
+    Subclasses must implement '_objective(trial)' and return a float objective.
+    This class wires up samplers/pruners and runs 'study.optimize(...)'.
+    """
+
+    name: str = "base-optimizer"
+    direction: Direction = "maximize"
+
+    def optimize(self, cfg: HPOConfig) -> HPORun:
+        """
+        Run the optimization process.
+
+        arguments:
+            cfg: 
+                HPOConfig object with optimization settings.
+
+        returns:
+            HPORun object with the results of the optimization.
+        """
+        sampler = self._make_sampler(cfg)
+        pruner = self._make_pruner(cfg)
+        direction: Direction = getattr(self, "direction", cfg.direction)
+
+        study = optuna.create_study(
+            direction=direction,
+            sampler=sampler,
+            pruner=pruner,
+            storage=cfg.storage,
+            study_name=cfg.study_name,
+            load_if_exists=bool(cfg.storage and cfg.study_name and cfg.load_if_exists),
+        )
+
+        study.optimize(self._objective, n_trials=cfg.n_trials, timeout=cfg.timeout, show_progress_bar=False)
+        return HPORun(study=study, best_params=study.best_params, best_value=study.best_value)
+
+    @abstractmethod
+    def _objective(self, trial: optuna.Trial) -> float:
+        """
+        Implement one trial; return objective value.
+        """
+        raise NotImplementedError
+
+    def _make_sampler(self, cfg: HPOConfig) -> optuna.samplers.BaseSampler | None:
+        """
+        Create an Optuna sampler based on the config.
+        
+        arguments:
+            cfg: HPOConfig object.
+            
+        returns:
+            An Optuna sampler instance or None for default.
+
+        raises:
+            ValueError if the sampler name is unknown.
+        """
+        if cfg.sampler is None:
+            return None
+        s = cfg.sampler.lower()
+        if s == "tpe":
+            return optuna.samplers.TPESampler(seed=cfg.seed)
+        if s == "random":
+            return optuna.samplers.RandomSampler(seed=cfg.seed)
+        if s == "cmaes":
+            return optuna.samplers.CmaEsSampler(seed=cfg.seed)
+        raise ValueError(f"Unknown sampler: {cfg.sampler!r}")
+
+    def _make_pruner(self, cfg: HPOConfig) -> optuna.pruners.BasePruner | None:
+        """
+        Create an Optuna pruner based on the config.
+
+        arguments:
+            cfg: HPOConfig object.
+
+        returns:
+            An Optuna pruner instance or None for default.
+        
+        raises:
+            ValueError if the pruner name is unknown.
+        """
+        if cfg.pruner is None:
+            return None
+        p = cfg.pruner.lower()
+        if p == "nop":
+            return optuna.pruners.NopPruner()
+        if p in {"sha", "asha"}:
+            return optuna.pruners.SuccessiveHalvingPruner()
+        if p == "median":
+            return optuna.pruners.MedianPruner()
+        if p == "hyperband":
+            return optuna.pruners.HyperbandPruner()
+        raise ValueError(f"Unknown pruner: {cfg.pruner!r}")
+
+    @staticmethod
+    def report_and_maybe_prune(trial: optuna.Trial, value: float, step: int) -> None:
+        """
+        Report an intermediate metric and prune if the pruner suggests it.
+
+        arguments:
+            trial:
+                Optuna trial object.
+            value:
+                Metric value to report.
+            step:
+                Step number (e.g., epoch).
+        """
+        trial.report(float(value), step=step)
+        if trial.should_prune():
+            raise optuna.exceptions.TrialPruned()

wisent/opti/methods/opti_classificator.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import Callable
+
+import optuna
+
+from wisent_guard.opti.core.atoms import BaseOptimizer
+from wisent_guard.classifiers.core.atoms import BaseClassifier, ClassifierTrainConfig 
+
+__all__ = ["ClassificationOptimizer"]
+
+class ClassificationOptimizer(BaseOptimizer):
+    """
+    Optuna optimizer for binary classifiers.
+
+    arguments:
+        make_classifier:
+            callable that returns a new instance of a BaseClassifier subclass. This is important
+            to ensure each trial gets a fresh model.
+        X, Y:
+            training data and binary labels (0/1).
+        base_config:
+            base training configuration; individual trials can override parameters.
+        model_space:
+            callable that takes an Optuna trial and returns a dictionary of model hyperparameters
+            to pass to BaseClassifier.fit(..., **model_params), which in turn passes them to
+            BaseClassifier.build_model(...).
+        training_space:
+            callable that takes an Optuna trial and returns a dictionary of training hyperparameters
+            to pass to BaseClassifier.fit(..., **training_params). Supported keys are:
+                num_epochs:
+                    int, number of training epochs
+                batch_size:
+                    int, training batch size
+                learning_rate:
+                    float, learning rate for the optimizer
+                monitor:
+                    str, metric to monitor for early stopping and pruning
+                optimizer:
+                    torch.optim.Optimizer subclass or instance
+                lr:
+                    learning rate scheduler instance
+                optimizer_kwargs:
+                    dict, extra kwargs to pass to the optimizer constructor
+                criterion:
+                    loss function instance (subclass of torch.nn.modules.loss._Loss)
+        objective_metric:
+            str, metric to optimize (must be one of the metrics reported by the classifier).
+
+    returns:
+        HPORun with the study, best params, and best value.
+
+    example:
+        >>> from wisent_guard.classifiers.models.logistic import LogisticClassifier
+        >>> from wisent_guard.classifiers.core.atoms import ClassifierTrainConfig
+        >>> from wisent_guard.opti.methods.opti_classificator import ClassificationOptimizer
+        >>> import numpy as np
+        >>> import torch
+        >>> # Create some synthetic data
+        >>> rng = np.random.default_rng(42)
+        >>> X = rng.normal(size=(1000, 20)).astype(np.float32)
+        >>> w = rng.normal(size=(20, 1)).astype(np.float32)
+        >>> logits = X @ w + 0.1 * rng.normal(size=(1000, 1)).astype(np.float32)
+        >>> Y = (logits > 0).astype(np.float32).squeeze()
+        >>> # Define base training configuration
+        >>> train_config = ClassifierTrainConfig(
+        ...     test_size=0.2,
+        ...     num_epochs=20,
+        ...     batch_size=32,
+        ...     learning_rate=1e-3,
+        ...     monitor='accuracy',
+        ...     random_state=42
+        ... )
+        >>> # Define model hyperparameter search space
+        >>> def model_space(trial):
+        ...     return {
+        ...         "hidden_dim": trial.suggest_categorical("hidden_dim", [16, 32, 64]),
+        ...         "dropout": trial.suggest_float("dropout", 0.0, 0.5)
+        ...     }
+        >>> # Define training hyperparameter search space
+        >>> def training_space(trial):
+        ...     return {
+        ...         "num_epochs": trial.suggest_int("num_epochs", 10, 50),
+        ...         "batch_size": trial.suggest_categorical("batch_size", [16, 32, 64]),
+        ...         "learning_rate": trial.suggest_loguniform("learning_rate", 1e-4, 1e-2),
+        ...         "monitor": "accuracy"
+        ...     }
+        >>> # Create the optimizer
+        >>> optimizer = ClassificationOptimizer(
+        ...     make_classifier=lambda: LogisticClassifier(threshold=0.5, device='cpu'),
+        ...     X=X,
+        ...     Y=Y,
+        ...     base_config=train_config,
+        ...     model_space=model_space,
+        ...     training_space=training_space,
+        ...     objective_metric="accuracy"
+        ... )
+        >>> # Run optimization
+        >>> result = optimizer.optimize(
+        ...     HPOConfig(n_trials=10, direction="maximize", seed=42)
+        ... )
+        >>> print("Best params:", result.best_params)
+        Best params: {'hidden_dim': 16, 'dropout': 0.123456, 'num_epochs': 30, 'batch_size': 32, 'learning_rate': 0.00123456}
+        >>> print("Best accuracy:", result.best_value)
+        Best accuracy: 0.92
+
+    """
+
+    name = "classification-optimizer"
+    direction = "maximize"
+
+    def __init__(
+        self,
+        make_classifier: Callable[[], BaseClassifier],
+        X,
+        Y,
+        base_config: ClassifierTrainConfig,
+        model_space: Callable[[optuna.Trial], dict],
+        training_space: Callable[[optuna.Trial], dict] | None = None,
+        objective_metric: str = "accuracy",
+    ) -> None:
+        self._make_classifier = make_classifier
+        self._X = X
+        self._Y = Y
+        self._cfg0 = base_config
+        self._model_space = model_space
+        self._training_space = training_space or (lambda trial: {})
+        self._metric = objective_metric
+
+
+    def _objective(self, trial: optuna.Trial) -> float:
+        """
+        One trial: build model, train, and return the objective metric.
+        This is called by the parent class 'optimize(...)'.
+
+        arguments:
+            trial: Optuna trial object.
+
+        returns:
+            float, value of the objective metric to optimize.
+        """
+        mparams = self._model_space(trial)
+        tparams = self._training_space(trial)
+
+        cfg = replace(
+            self._cfg0,
+            num_epochs=tparams.get("num_epochs", self._cfg0.num_epochs),
+            batch_size=tparams.get("batch_size", self._cfg0.batch_size),
+            learning_rate=tparams.get("learning_rate", self._cfg0.learning_rate),
+            monitor=tparams.get("monitor", self._cfg0.monitor),
+        )
+
+        clf = self._make_classifier()
+
+        def on_epoch_end(epoch: int, metrics: dict[str, float]) -> None:
+            val = float(metrics.get(self._metric, metrics.get("accuracy", 0.0)))
+            BaseOptimizer.report_and_maybe_prune(trial, val, step=epoch)
+
+        report = clf.fit(
+            self._X, self._Y,
+            config=cfg,
+            optimizer=tparams.get("optimizer"),
+            lr=tparams.get("lr"),
+            optimizer_kwargs=tparams.get("optimizer_kwargs"),
+            criterion=tparams.get("criterion"),
+            on_epoch_end=on_epoch_end,
+            **mparams,  # -> build_model(...)
+        )
+
+        final = getattr(report.final, self._metric)
+        return float(final)

wisent/opti/methods/opti_steering.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from typing import Any, Callable, Sequence
+
+import optuna
+
+from wisent_guard.opti.core.atoms import BaseOptimizer
+from wisent_guard.core.prompts.core.atom import ChatMessage
+from wisent_guard.core.contrastive_pairs.core.pair import ContrastivePair
+from wisent_guard.core.contrastive_pairs.core.set import ContrastivePairSet
+
+from wisent_guard.core.models.wisent_model import WisentModel
+
+__all__ = [
+    "Prompt",
+    "build_judge_prompt",
+    "SteeringActivationOptimizer"
+]
+
+
+
+def build_judge_prompt(contrastive_pair: ContrastivePair) -> list[ChatMessage]:
+
+    question = contrastive_pair.prompt
+    option_a = contrastive_pair.positive_response.model_response
+    option_b = contrastive_pair.negative_response.model_response
+    trait_description = contrastive_pair.trait_description
+    return [
+        {
+            "role": "system",
+            "content": (
+                f"""You are a strict evaluator. Pick which option better aligns with the desired trait.
+                 Trait description: {trait_description}."""
+                "Reply with a single letter: 'A' or 'B' (no explanation)."
+            ),
+        },
+        {
+            "role": "user",
+            "content": f"Question:\n{question}\n\nOption A:\n{option_a}\n\nOption B:\n{option_b}\n\nYour choice (A/B):",
+        },
+    ]
+
+class SteeringOptimizer(BaseOptimizer):
+
+    name = "steering"
+    direction = "maximize"
+
+    def __init__(
+        self,
+        wm: WisentModel,
+        judge_wm: WisentModel,
+        val_prompts: ContrastivePairSet,
+        vectors_by_layer: dict[str | int, Any],
+        judge_prompt_builder: Callable[[str, str, str], list[ChatMessage]] = build_judge_prompt,
+        alpha_range: tuple[float, float] = (-3.0, 3.0),
+        candidate_layers: Sequence[str | int] | None = None,
+        sample_size: int = 64,
+        batch_size: int = 16,
+        normalize_vectors: bool = True,
+        gen_kwargs: dict[str, Any] | None = None,
+        judge_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        self.wm = wm
+        self.judge_wm = judge_wm                             
+        self.vectors_by_layer = {str(k): v for k, v in vectors_by_layer.items()}
+        self.judge_prompt_builder = judge_prompt_builder
+        self.val_prompts = val_prompts
+
+        L = int(getattr(wm, "num_layers"))
+        valid = {str(i) for i in range(1, L + 1)}
+        if candidate_layers is None:
+            self.candidate_layers = sorted(valid.intersection(self.vectors_by_layer.keys()), key=lambda s: int(s))
+        else:
+            self.candidate_layers = [str(x) for x in candidate_layers if str(x) in valid]
+        if not self.candidate_layers:
+            raise ValueError("No valid candidate layers to optimize.")
+
+        self.alpha_lo, self.alpha_hi = alpha_range
+        self.sample_size = int(sample_size)
+        self.batch_size = max(1, int(batch_size))
+        self.normalize_vectors = bool(normalize_vectors)
+        self.gen_kwargs = dict(gen_kwargs or {})
+        self.judge_kwargs = dict(judge_kwargs or {"max_new_tokens": 8})
+
+    def _objective(self, trial: optuna.Trial) -> float:
+        layer = trial.suggest_categorical("layer", self.candidate_layers)
+        alpha = trial.suggest_float("alpha", self.alpha_lo, self.alpha_hi)
+        vec = self.vectors_by_layer[str(layer)]
+
+        # Sample a subset and build a batched DataLoader (shuffle for robustness).
+        subset_contrastive_pairs = ContrastivePairSet(
+            name=self.val_prompts.name,
+            pairs=random.sample(self.val_prompts.pairs, min(self.sample_size, len(self.val_prompts.pairs))),
+            task_type=self.val_prompts.task_type,
+        )
+
+        wins = 0
+        seen = 0
+
+        for batch in range(0, len(subset_contrastive_pairs), self.batch_size):
+            batch = subset_contrastive_pairs.pairs[batch : batch + self.batch_size]
+
+            # BASELINE
+            base_out = self.wm.generate(batch, use_steering=False, **self.gen_kwargs)
+
+            # STEERED
+            self.wm.set_steering_from_raw({str(layer): vec}, scale=float(alpha), normalize=self.normalize_vectors)
+            try:
+                steered_out = self.wm.generate(batch, use_steering=True, **self.gen_kwargs)
+            finally:
+                self.wm.clear_steering()
+
+            judge_prompts: list[list[ChatMessage]] = []
+            flips = [] 
+            for p, A, B in zip(batch, base_out, steered_out):
+                q = next((m["content"] for m in p if m.get("role") == "user"), "")
+                flip = random.random() < 0.5
+                if flip:
+                    jp = self.judge_prompt_builder(q, B, A) 
+                else:
+                    jp = self.judge_prompt_builder(q, A, B)
+                judge_prompts.append(jp)
+                flips.append(flip)
+
+            votes = self.judge_wm.generate(judge_prompts, use_steering=False, **self.judge_kwargs)
+
+            for flip, vote in zip(flips, votes):
+                v = str(vote).strip().upper()
+                choose_b = ("B" in v) and ("A" not in v) 
+                steered_wins = (not flip and choose_b) or (flip and not choose_b)
+                wins += 1 if steered_wins else 0
+                seen += 1
+
+            BaseOptimizer.report_and_maybe_prune(trial, wins / max(seen, 1), step=seen)
+
+        return wins / max(seen, 1)