wisent 0.5.12__py3-none-any.whl → 0.5.13__py3-none-any.whl

wisent/core/parser_arguments/generate_vector_parser.py

@@ -0,0 +1,90 @@
+"""Parser setup for the 'generate-vector' command."""
+
+
+def setup_generate_vector_parser(parser):
+    """Set up the generate-vector subcommand parser."""
+    # Source of contrastive pairs - mutually exclusive for single property
+    source_group = parser.add_mutually_exclusive_group(required=False)
+    source_group.add_argument(
+        "--from-pairs",
+        type=str,
+        metavar="FILE",
+        help="Path to JSON file containing contrastive pairs (single property)",
+    )
+    source_group.add_argument(
+        "--from-description",
+        type=str,
+        metavar="TRAIT",
+        help="Natural language description of the trait (single property)",
+    )
+
+    # Multi-property support
+    parser.add_argument("--multi-property", action="store_true", help="Enable multi-property steering (DAC only)")
+    parser.add_argument(
+        "--property-files",
+        type=str,
+        nargs="+",
+        metavar="NAME:FILE:LAYER",
+        help="Property definitions from files (format: property_name:pairs_file:layer)",
+    )
+    parser.add_argument(
+        "--property-descriptions",
+        type=str,
+        nargs="+",
+        metavar="NAME:DESC:LAYER",
+        help="Property definitions from descriptions (format: property_name:description:layer)",
+    )
+
+    # Model configuration
+    parser.add_argument("--model", type=str, default="distilgpt2", help="Model name or path (default: distilgpt2)")
+    parser.add_argument("--device", type=str, default=None, help="Device to run on (default: auto-detect)")
+
+    # Steering method configuration
+    parser.add_argument(
+        "--method",
+        type=str,
+        default="DAC",
+        choices=["DAC", "CAA", "HPR", "BiPO", "ControlVectorSteering"],
+        help="Steering method to use (default: DAC)",
+    )
+    parser.add_argument("--layer", type=int, default=0, help="Layer index to apply steering (default: 0)")
+
+    # Output configuration
+    parser.add_argument("--output", type=str, required=True, help="Output path for the generated steering vector")
+
+    # Pair generation options (only used with --from-description)
+    parser.add_argument(
+        "--num-pairs",
+        type=int,
+        default=30,
+        help="Number of pairs to generate when using --from-description (default: 30)",
+    )
+    parser.add_argument(
+        "--save-pairs", type=str, default=None, help="Save generated pairs to this file when using --from-description"
+    )
+
+    # Method-specific parameters
+    parser.add_argument("--dynamic-control", action="store_true", help="Enable dynamic control for DAC method")
+    parser.add_argument(
+        "--entropy-threshold", type=float, default=1.0, help="Entropy threshold for DAC method (default: 1.0)"
+    )
+    parser.add_argument("--beta", type=float, default=1.0, help="Beta parameter for HPR method (default: 1.0)")
+
+    # Activation extraction configuration
+    parser.add_argument(
+        "--prompt-construction",
+        type=str,
+        default="multiple_choice",
+        choices=["multiple_choice", "role_playing", "direct_completion", "instruction_following"],
+        help="Strategy for constructing prompts from question-answer pairs (default: multiple_choice)",
+    )
+    parser.add_argument(
+        "--token-targeting",
+        type=str,
+        default="choice_token",
+        choices=["choice_token", "continuation_token", "last_token", "first_token", "mean_pooling", "max_pooling"],
+        help="Strategy for targeting tokens in activation extraction (default: choice_token)",
+    )
+
+    # General options
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")

wisent/core/parser_arguments/get_activations_parser.py

@@ -0,0 +1,90 @@
+"""Parser for the get-activations command."""
+
+import argparse
+
+
+def setup_get_activations_parser(parser: argparse.ArgumentParser) -> None:
+    """
+    Set up the get-activations command parser.
+
+    This command loads contrastive pairs from a JSON file, collects activations
+    from specified model layers, and saves the enriched pairs back to disk.
+    """
+    # Input/Output
+    parser.add_argument(
+        "pairs_file",
+        type=str,
+        help="Path to JSON file containing contrastive pairs"
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        required=True,
+        help="Output file path for pairs with activations (JSON)"
+    )
+
+    # Model configuration
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="meta-llama/Llama-3.2-1B-Instruct",
+        help="Model identifier (e.g., 'meta-llama/Llama-3.2-1B-Instruct')"
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="Device to run on (cuda, cpu, mps)"
+    )
+
+    # Layer selection
+    parser.add_argument(
+        "--layers",
+        type=str,
+        default=None,
+        help="Comma-separated layer indices (e.g., '8,12,15') or 'all' for all layers"
+    )
+
+    # Token aggregation
+    parser.add_argument(
+        "--token-aggregation",
+        type=str,
+        choices=["average", "final", "first", "max", "min"],
+        default="average",
+        help="How to aggregate token activations"
+    )
+
+    # Prompt construction strategy
+    parser.add_argument(
+        "--prompt-strategy",
+        type=str,
+        choices=["chat_template", "direct_completion", "instruction_following", "multiple_choice", "role_playing"],
+        default="chat_template",
+        help="Prompt construction strategy (default: chat_template)"
+    )
+
+    # Processing options
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        help="Batch size for processing (default: 1)"
+    )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="Maximum number of pairs to process"
+    )
+
+    # Display options
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="Enable verbose output"
+    )
+    parser.add_argument(
+        "--timing",
+        action="store_true",
+        help="Show timing information"
+    )

wisent/core/parser_arguments/main_parser.py

@@ -0,0 +1,152 @@
+"""
+Main parser setup for Wisent CLI.
+
+This module imports and combines all command-specific parsers into a single
+argparse parser for the Wisent CLI.
+"""
+
+import argparse
+
+from wisent.core.parser_arguments.tasks_parser import setup_tasks_parser
+from wisent.core.parser_arguments.generate_pairs_parser import setup_generate_pairs_parser
+from wisent.core.parser_arguments.generate_pairs_from_task_parser import setup_generate_pairs_from_task_parser
+from wisent.core.parser_arguments.get_activations_parser import setup_get_activations_parser
+from wisent.core.parser_arguments.create_steering_vector_parser import setup_create_steering_vector_parser
+from wisent.core.parser_arguments.generate_vector_from_task_parser import setup_generate_vector_from_task_parser
+from wisent.core.parser_arguments.generate_vector_from_synthetic_parser import setup_generate_vector_from_synthetic_parser
+from wisent.core.parser_arguments.synthetic_parser import setup_synthetic_parser
+from wisent.core.parser_arguments.test_nonsense_parser import setup_test_nonsense_parser
+from wisent.core.parser_arguments.monitor_parser import setup_monitor_parser
+from wisent.core.parser_arguments.agent_parser import setup_agent_parser
+from wisent.core.parser_arguments.model_config_parser import setup_model_config_parser
+from wisent.core.parser_arguments.configure_model_parser import setup_configure_model_parser
+from wisent.core.parser_arguments.optimize_classification_parser import setup_classification_optimizer_parser
+from wisent.core.parser_arguments.optimize_steering_parser import setup_steering_optimizer_parser
+from wisent.core.parser_arguments.optimize_sample_size_parser import setup_sample_size_optimizer_parser
+from wisent.core.parser_arguments.full_optimize_parser import setup_full_optimizer_parser
+from wisent.core.parser_arguments.generate_vector_parser import setup_generate_vector_parser
+from wisent.core.parser_arguments.multi_steer_parser import setup_multi_steer_parser
+from wisent.core.parser_arguments.evaluate_parser import setup_evaluate_parser
+from wisent.core.parser_arguments.generate_responses_parser import setup_generate_responses_parser
+from wisent.core.parser_arguments.evaluate_responses_parser import setup_evaluate_responses_parser
+
+
+def setup_parser() -> argparse.ArgumentParser:
+    """Set up the main CLI parser with subcommands."""
+    parser = argparse.ArgumentParser(description="Wisent-Guard: Advanced AI Safety and Alignment Toolkit")
+
+    # Global arguments
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
+
+    # Create subparsers for different commands
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+
+    # Tasks command (main evaluation pipeline)
+    tasks_parser = subparsers.add_parser("tasks", help="Run evaluation tasks")
+    setup_tasks_parser(tasks_parser)
+
+    # Generate pairs command
+    generate_parser = subparsers.add_parser("generate-pairs", help="Generate synthetic contrastive pairs")
+    setup_generate_pairs_parser(generate_parser)
+
+    # Generate pairs from task command
+    generate_from_task_parser = subparsers.add_parser("generate-pairs-from-task", help="Generate contrastive pairs from lm-eval task")
+
+    # Get activations command
+    get_activations_parser = subparsers.add_parser("get-activations", help="Collect activations from contrastive pairs")
+    setup_get_activations_parser(get_activations_parser)
+    setup_generate_pairs_from_task_parser(generate_from_task_parser)
+
+    # Create steering vector command
+    create_steering_parser = subparsers.add_parser("create-steering-vector", help="Create steering vectors from enriched pairs")
+    setup_create_steering_vector_parser(create_steering_parser)
+
+    # Generate vector from task command (full pipeline)
+    generate_vector_from_task_parser = subparsers.add_parser("generate-vector-from-task", help="Generate steering vector from task (full pipeline)")
+    setup_generate_vector_from_task_parser(generate_vector_from_task_parser)
+
+    # Generate vector from synthetic command (full pipeline)
+    generate_vector_from_synthetic_parser = subparsers.add_parser("generate-vector-from-synthetic", help="Generate steering vector from synthetic pairs (full pipeline)")
+    setup_generate_vector_from_synthetic_parser(generate_vector_from_synthetic_parser)
+
+    # Synthetic command (generate + train + test)
+    synthetic_parser = subparsers.add_parser("synthetic", help="Run synthetic contrastive pair pipeline")
+    setup_synthetic_parser(synthetic_parser)
+
+    # Test nonsense detection command
+    test_nonsense_parser = subparsers.add_parser("test-nonsense", help="Test nonsense detection system")
+    setup_test_nonsense_parser(test_nonsense_parser)
+
+    # Monitor command for performance monitoring
+    monitor_parser = subparsers.add_parser("monitor", help="Performance monitoring and system information")
+    setup_monitor_parser(monitor_parser)
+
+    # Agent command for autonomous agent interaction
+    agent_parser = subparsers.add_parser("agent", help="Interact with autonomous agent")
+    setup_agent_parser(agent_parser)
+
+    # Model configuration command for managing optimal parameters
+    model_config_parser = subparsers.add_parser("model-config", help="Manage model-specific optimal parameters")
+    setup_model_config_parser(model_config_parser)
+
+    # Configure model command for setting up new/unsupported models
+    configure_model_parser = subparsers.add_parser(
+        "configure-model", help="Configure tokens and layer access for unsupported models"
+    )
+    setup_configure_model_parser(configure_model_parser)
+
+    # Classification optimization command for finding optimal classification parameters
+    classification_optimizer_parser = subparsers.add_parser(
+        "optimize-classification", help="Optimize classification parameters across all tasks"
+    )
+    setup_classification_optimizer_parser(classification_optimizer_parser)
+
+    # Steering optimization command for finding optimal steering parameters
+    steering_optimizer_parser = subparsers.add_parser(
+        "optimize-steering", help="Optimize steering parameters for different methods"
+    )
+    setup_steering_optimizer_parser(steering_optimizer_parser)
+
+    # Sample size optimization command for finding optimal training sample sizes
+    sample_size_optimizer_parser = subparsers.add_parser(
+        "optimize-sample-size", help="Find optimal training sample size for classifiers"
+    )
+    setup_sample_size_optimizer_parser(sample_size_optimizer_parser)
+
+    # Full optimization command that runs both classification and sample size optimization
+    full_optimizer_parser = subparsers.add_parser(
+        "full-optimize", help="Run full optimization: classification parameters then sample size"
+    )
+    setup_full_optimizer_parser(full_optimizer_parser)
+
+    # Generate vector command for creating steering vectors without tasks
+    generate_vector_parser = subparsers.add_parser(
+        "generate-vector", help="Generate steering vectors from contrastive pairs (file or description)"
+    )
+    setup_generate_vector_parser(generate_vector_parser)
+
+    # Multi-vector steering command for combining multiple vectors at inference time
+    multi_steer_parser = subparsers.add_parser(
+        "multi-steer", help="Combine multiple steering vectors dynamically at inference time"
+    )
+    setup_multi_steer_parser(multi_steer_parser)
+
+    # Single-prompt evaluation command for real-time steering assessment
+    evaluate_parser = subparsers.add_parser(
+        "evaluate", help="Evaluate single prompt with steering vector and return quality scores"
+    )
+    setup_evaluate_parser(evaluate_parser)
+
+    # Generate responses command for generating model responses to task questions
+    generate_responses_parser = subparsers.add_parser(
+        "generate-responses", help="Generate model responses to questions from a task"
+    )
+    setup_generate_responses_parser(generate_responses_parser)
+
+    # Evaluate responses command for evaluating generated responses
+    evaluate_responses_parser = subparsers.add_parser(
+        "evaluate-responses", help="Evaluate generated responses using embedded evaluator"
+    )
+    setup_evaluate_responses_parser(evaluate_responses_parser)
+
+    return parser

wisent/core/parser_arguments/model_config_parser.py

@@ -0,0 +1,59 @@
+"""Parser setup for the 'model-config' command."""
+
+
+def setup_model_config_parser(parser):
+    """Set up the model-config subcommand parser."""
+    # Create subparsers for different model config actions
+    config_subparsers = parser.add_subparsers(dest="config_action", help="Model configuration actions")
+
+    # Save configuration subcommand
+    save_parser = config_subparsers.add_parser("save", help="Save optimal parameters for a model")
+    save_parser.add_argument("model", type=str, help="Model name or path")
+    save_parser.add_argument("--classification-layer", type=int, required=True, help="Optimal layer for classification")
+    save_parser.add_argument(
+        "--steering-layer", type=int, default=None, help="Optimal layer for steering (defaults to classification layer)"
+    )
+    save_parser.add_argument(
+        "--token-aggregation",
+        type=str,
+        default="average",
+        choices=["average", "final", "first", "max", "min"],
+        help="Token aggregation method",
+    )
+    save_parser.add_argument("--detection-threshold", type=float, default=0.6, help="Detection threshold")
+    save_parser.add_argument(
+        "--optimization-method", type=str, default="manual", help="How these parameters were determined"
+    )
+    save_parser.add_argument("--metrics", type=str, default=None, help="JSON string with optimization metrics")
+
+    # List configurations subcommand
+    list_parser = config_subparsers.add_parser("list", help="List all saved model configurations")
+    list_parser.add_argument("--detailed", action="store_true", help="Show detailed configuration information")
+
+    # Show configuration subcommand
+    show_parser = config_subparsers.add_parser("show", help="Show configuration for a specific model")
+    show_parser.add_argument("model", type=str, help="Model name or path")
+    show_parser.add_argument("--task", type=str, default=None, help="Show task-specific overrides if available")
+
+    # Remove configuration subcommand
+    remove_parser = config_subparsers.add_parser("remove", help="Remove configuration for a model")
+    remove_parser.add_argument("model", type=str, help="Model name or path")
+    remove_parser.add_argument("--confirm", action="store_true", help="Confirm removal without prompting")
+
+    # Test configuration subcommand
+    test_parser = config_subparsers.add_parser("test", help="Test if saved configuration works")
+    test_parser.add_argument("model", type=str, help="Model name or path")
+    test_parser.add_argument(
+        "--task", type=str, default="truthfulqa_mc1", help="Task to test with (default: truthfulqa_mc1)"
+    )
+    test_parser.add_argument("--limit", type=int, default=5, help="Number of samples to test with (default: 5)")
+    test_parser.add_argument("--device", type=str, default=None, help="Device to run on")
+
+    # Common arguments for all subcommands
+    parser.add_argument(
+        "--config-dir",
+        type=str,
+        default=None,
+        help="Custom directory for configuration files (default: ~/.wisent-guard/model_configs/)",
+    )
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")

wisent/core/parser_arguments/monitor_parser.py

@@ -0,0 +1,17 @@
+"""Parser setup for the 'monitor' command."""
+
+
+def setup_monitor_parser(parser):
+    """Set up the monitor subcommand parser."""
+    parser.add_argument("--memory-info", action="store_true", help="Show current memory usage information")
+    parser.add_argument("--system-info", action="store_true", help="Show system information and capabilities")
+    parser.add_argument("--benchmark", action="store_true", help="Run performance benchmarks")
+    parser.add_argument("--test-gpu", action="store_true", help="Test GPU availability and memory")
+    parser.add_argument("--continuous", action="store_true", help="Continuous monitoring mode (Ctrl+C to stop)")
+    parser.add_argument("--interval", type=float, default=1.0, help="Monitoring interval in seconds (default: 1.0)")
+    parser.add_argument("--export-csv", type=str, default=None, help="Export monitoring data to CSV file")
+    parser.add_argument(
+        "--duration", type=int, default=60, help="Duration for continuous monitoring in seconds (default: 60)"
+    )
+    parser.add_argument("--track-gpu", action="store_true", help="Include GPU monitoring (requires CUDA)")
+    parser.add_argument("--detailed", action="store_true", help="Show detailed monitoring information")

wisent/core/parser_arguments/multi_steer_parser.py

@@ -0,0 +1,47 @@
+"""Parser setup for the 'multi-steer' command."""
+
+
+def setup_multi_steer_parser(parser):
+    """Set up the multi-steer subcommand parser for dynamic vector combination."""
+    # Vector inputs - can specify multiple vector-weight pairs
+    parser.add_argument(
+        "--vector",
+        type=str,
+        action="append",
+        required=True,
+        metavar="PATH:WEIGHT",
+        help="Path to steering vector and its weight (format: path/to/vector.pt:0.5). Can be specified multiple times.",
+    )
+
+    # Model configuration
+    parser.add_argument("--model", type=str, required=True, help="Model name or path")
+    parser.add_argument("--layer", type=int, required=True, help="Layer index to apply combined steering")
+    parser.add_argument("--device", type=str, default=None, help="Device to run on (default: auto-detect)")
+
+    # Steering method configuration
+    parser.add_argument(
+        "--method",
+        type=str,
+        default="CAA",
+        choices=["CAA", "DAC"],
+        help="Steering method to use for combination (default: CAA)",
+    )
+
+    # Generation configuration
+    parser.add_argument("--prompt", type=str, required=True, help="Prompt to generate with combined steering")
+    parser.add_argument("--max-new-tokens", type=int, default=100, help="Maximum new tokens to generate (default: 100)")
+
+    # Weight normalization
+    parser.add_argument("--normalize-weights", action="store_true", help="Normalize weights to sum to 1.0")
+    parser.add_argument(
+        "--allow-unnormalized", action="store_true", help="Allow weights that don't sum to 1.0 (for stronger effects)"
+    )
+    parser.add_argument(
+        "--target-norm", type=float, default=None, help="Scale the combined vector to have this norm (e.g., 10.0)"
+    )
+
+    # Output options
+    parser.add_argument(
+        "--save-combined", type=str, default=None, help="Save the combined steering vector to this path"
+    )
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose output showing weight calculations")

wisent/core/parser_arguments/optimize_classification_parser.py

@@ -0,0 +1,67 @@
+"""Parser setup for the 'optimize-classification' command."""
+
+
+def setup_classification_optimizer_parser(parser):
+    """Set up the classification-optimizer subcommand parser."""
+    parser.add_argument("model", type=str, help="Model name or path to optimize")
+    parser.add_argument("--limit", type=int, default=1000, help="Maximum samples per task (default: 1000)")
+    parser.add_argument(
+        "--optimization-metric",
+        type=str,
+        default="f1",
+        choices=["f1", "accuracy", "precision", "recall"],
+        help="Metric to optimize (default: f1)",
+    )
+    parser.add_argument(
+        "--max-time-per-task", type=float, default=15.0, help="Maximum time per task in minutes (default: 15.0)"
+    )
+    parser.add_argument(
+        "--layer-range", type=str, default=None, help="Layer range to test (e.g., '10-20', if None uses all layers)"
+    )
+    parser.add_argument(
+        "--aggregation-methods",
+        type=str,
+        nargs="+",
+        default=["average", "final", "first", "max", "min"],
+        help="Token aggregation methods to test",
+    )
+    parser.add_argument(
+        "--threshold-range",
+        type=float,
+        nargs="+",
+        default=[0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
+        help="Detection thresholds to test",
+    )
+    parser.add_argument("--device", type=str, default=None, help="Device to run on")
+    parser.add_argument("--results-file", type=str, default=None, help="Custom file path for saving results")
+    parser.add_argument("--no-save", action="store_true", help="Don't save results to model config")
+    parser.add_argument("--save-logs-json", type=str, default=None, help="Save detailed optimization logs to JSON file")
+    parser.add_argument(
+        "--save-classifiers",
+        action="store_true",
+        default=True,
+        help="Save best classifiers for each task (default: True)",
+    )
+    parser.add_argument(
+        "--no-save-classifiers",
+        dest="save_classifiers",
+        action="store_false",
+        help="Don't save classifiers (overrides --save-classifiers)",
+    )
+    parser.add_argument(
+        "--classifiers-dir",
+        type=str,
+        default=None,
+        help="Directory to save classifiers (default: ./optimized_classifiers/model_name/)",
+    )
+
+    # Timing calibration options
+    parser.add_argument(
+        "--skip-timing-estimation", action="store_true", help="Skip timing estimation and proceed without time warnings"
+    )
+    parser.add_argument("--calibration-file", type=str, default=None, help="File to save/load calibration data")
+    parser.add_argument(
+        "--calibrate-only",
+        action="store_true",
+        help="Only run calibration and exit (saves to --calibration-file if provided)",
+    )

wisent/core/parser_arguments/optimize_sample_size_parser.py

@@ -0,0 +1,58 @@
+"""Parser setup for the 'optimize-sample-size' command."""
+
+
+def setup_sample_size_optimizer_parser(parser):
+    """Set up the sample-size-optimizer subcommand parser."""
+    parser.add_argument("model", type=str, help="Model name or path to optimize")
+    parser.add_argument("--task", type=str, required=True, help="Task to optimize for (REQUIRED)")
+    parser.add_argument("--layer", type=int, required=True, help="Layer index to use (REQUIRED)")
+    parser.add_argument(
+        "--token-aggregation",
+        type=str,
+        required=True,
+        choices=["average", "final", "first", "max", "min"],
+        help="Token aggregation method (REQUIRED)",
+    )
+
+    # Classification-specific arguments
+    parser.add_argument(
+        "--threshold", type=float, default=0.5, help="Detection threshold for classification (default: 0.5)"
+    )
+
+    # Steering mode
+    parser.add_argument("--steering-mode", action="store_true", help="Optimize for steering instead of classification")
+    parser.add_argument(
+        "--steering-method",
+        type=str,
+        default="CAA",
+        choices=["CAA", "CAA_L2", "HPR", "DAC", "BiPO", "KSteering"],
+        help="Steering method to use (default: CAA)",
+    )
+    parser.add_argument("--steering-strength", type=float, default=1.0, help="Steering strength to use (default: 1.0)")
+    parser.add_argument(
+        "--token-targeting-strategy",
+        type=str,
+        default="LAST_TOKEN",
+        choices=["CHOICE_TOKEN", "LAST_TOKEN", "FIRST_TOKEN", "ALL_TOKENS"],
+        help="Token targeting strategy for steering (default: LAST_TOKEN)",
+    )
+
+    # Common optimization parameters
+    parser.add_argument(
+        "--sample-sizes",
+        type=int,
+        nargs="+",
+        default=[5, 10, 20, 50, 100, 200, 500],
+        help="Sample sizes to test (default: 5 10 20 50 100 200 500)",
+    )
+    parser.add_argument("--test-size", type=int, default=200, help="Fixed test set size (default: 200)")
+    parser.add_argument("--test-split", type=float, default=0.2, help="DEPRECATED: Use --test-size instead")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility (default: 42)")
+    parser.add_argument("--limit", type=int, default=None, help="Maximum number of samples to load from dataset")
+    parser.add_argument("--save-plot", action="store_true", help="Save performance plot")
+    parser.add_argument("--no-save-config", action="store_true", help="Don't save optimal sample size to model config")
+    parser.add_argument("--device", type=str, default=None, help="Device to run on")
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
+    parser.add_argument(
+        "--force", action="store_true", help="Force optimization even without matching classifier parameters"
+    )