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

wisent/cli/wisent_cli/commands/train_cmd.py

@@ -1,322 +0,0 @@
-from __future__ import annotations
-import json
-import os
-from pathlib import Path
-from typing import Dict, List, Optional
-
-import typer
-
-from wisent.cli.wisent_cli.ui import echo
-from wisent.cli.wisent_cli.util import aggregations as aggs
-from wisent.cli.wisent_cli.util.parsing import (
-    parse_natural_tokens, parse_kv, parse_layers, to_bool, DTYPE_MAP,
-)
-
-try:
-    from rich.table import Table
-    from rich.panel import Panel
-    from rich.syntax import Syntax
-    HAS_RICH = True
-except Exception:
-    HAS_RICH = False
-
-__all__ = ["app", "train"]
-
-app = typer.Typer(help="Training workflow")
-
-def _resolve_method(method_name: Optional[str], methods_location: Optional[str]):
-    from wisent.cli.steering_methods.steering_rotator import SteeringMethodRotator  # type: ignore
-    # Best effort discovery if available
-    try:
-        if methods_location and hasattr(SteeringMethodRotator, "discover_methods"):
-            SteeringMethodRotator.discover_methods(methods_location)  # type: ignore[attr-defined]
-    except Exception:
-        pass
-
-    rot = SteeringMethodRotator()
-    if method_name:
-        # Case-insensitive match from registry
-        registry = {m["name"].lower(): m["name"] for m in rot.list_methods()}
-        real = registry.get(method_name.lower(), method_name)
-        try:
-            rot.use(real)
-            inst = getattr(rot, "_method", None)
-            if inst is not None:
-                return inst
-        except Exception:
-            pass
-        # Fallback to private resolver
-        try:
-            return SteeringMethodRotator._resolve_method(real)
-        except Exception as ex:
-            raise typer.BadParameter(f"Unknown steering method: {method_name!r}") from ex
-
-    # No name provided -> default to first or 'caa' if present
-    names = [m["name"] for m in rot.list_methods()]
-    if "caa" in [n.lower() for n in names]:
-        rot.use("caa")
-        return getattr(rot, "_method", SteeringMethodRotator._resolve_method("caa"))
-    if not names:
-        raise typer.BadParameter("No steering methods registered.")
-    rot.use(names[0])
-    return getattr(rot, "_method", SteeringMethodRotator._resolve_method(names[0]))
-
-
-def _show_plan(
-    *,
-    model: str,
-    loader: Optional[str],
-    loaders_location: Optional[str],
-    loader_kwargs: Dict[str, object],
-    method_name: Optional[str],
-    method_kwargs: Dict[str, object],
-    layers: Optional[str],
-    aggregation_name: str,
-    store_device: str,
-    dtype: Optional[str],
-    return_full_sequence: bool,
-    normalize_layers: bool,
-    save_dir: Optional[Path],
-) -> None:
-    plan = {
-        "Model": model,
-        "Data loader": loader or "(default)",
-        "Loaders location": loaders_location or "(auto)",
-        "Loader kwargs": loader_kwargs or {},
-        "Method": method_name or "(resolved automatically)",
-        "Method kwargs": method_kwargs or {},
-        "Layers": layers or "(all)",
-        "Aggregation": aggregation_name,
-        "Return full sequence": return_full_sequence,
-        "Normalize layers": normalize_layers,
-        "Store device": store_device,
-        "Dtype": dtype or "(unchanged)",
-        "Save dir": str(save_dir) if save_dir else "(none)",
-    }
-
-    code = f"""
-# Example: Training steering vectors (auto-generated plan)
-from wisent.core.trainers.steering_trainer import WisentSteeringTrainer
-from wisent.core.models.wisent_model import WisentModel
-from wisent.cli.data_loaders.data_loader_rotator import DataLoaderRotator
-from wisent.cli.steering_methods.steering_rotator import SteeringMethodRotator
-from wisent.core.activations.core.atoms import ActivationAggregationStrategy
-
-# 1) Model
-model = WisentModel(model_name={model!r}, layers={{}}, device={store_device!r})
-
-# 2) Data loader
-rot = DataLoaderRotator(loader={loader!r}, loaders_location={loaders_location!r})
-load = rot.load(**{json.dumps(loader_kwargs)})
-
-# 3) Method
-method = SteeringMethodRotator._resolve_method({(method_name or 'caa')!r})
-
-# 4) Trainer
-trainer = WisentSteeringTrainer(model=model, pair_set=load["train_qa_pairs"], steering_method=method,
-                                store_device={store_device!r}, dtype={dtype!r})
-
-# 5) Train
-result = trainer.run(
-    layers_spec={layers!r},
-    method_kwargs={json.dumps(method_kwargs)},
-    aggregation=ActivationAggregationStrategy.{aggs.pick(aggregation_name).name},
-    return_full_sequence={return_full_sequence!r},
-    normalize_layers={normalize_layers!r},
-    save_dir={str(save_dir) if save_dir else None!r},
-)
-""".strip()
-
-    if HAS_RICH:
-        t = Table(title="Execution Plan")
-        t.add_column("Key", style="bold", no_wrap=True)
-        t.add_column("Value")
-        for k, v in plan.items():
-            t.add_row(k, json.dumps(v) if isinstance(v, (dict, list)) else str(v))
-        echo(Panel(t, expand=False))
-        echo(Panel(Syntax(code, "python", word_wrap=False), title="Code Preview", expand=False))
-    else:
-        print(json.dumps(plan, indent=2))
-        print("\n" + code)
-
-
-@app.command("train", context_settings={"ignore_unknown_options": True, "allow_extra_args": True})
-def train(ctx: typer.Context, params: List[str] = typer.Argument(None)):
-    """
-    Natural (no-dash) usage examples:
-
-      wisent train model meta-llama/Llama-3.2-1B-Instruct loader custom path ./custom.json training_limit 5 method caa
-
-      wisent train interactive true
-
-    See `wisent loader-args custom` to view the exact loader arguments.
-    """
-    # Lazy imports
-    from wisent.cli.data_loaders.data_loader_rotator import DataLoaderRotator  # type: ignore
-    from wisent.core.models.wisent_model import WisentModel  # type: ignore
-    from wisent.core.trainers.steering_trainer import WisentSteeringTrainer  # type: ignore
-
-    tokens = list(params or []) + list(ctx.args or [])
-    top, loader_kv_raw, method_kv_raw = parse_natural_tokens(tokens)
-
-    # Core args
-    model = top.get("model")
-    if not model:
-        raise typer.BadParameter("Please specify a model (e.g. `train model meta-llama/Llama-3.2-1B-Instruct`) or use `interactive true`.")
-
-    loader = top.get("loader")
-    loaders_location = top.get("loaders_location")
-    methods_location = top.get("methods_location")
-    method_name = top.get("method")
-
-    layers = parse_layers(top.get("layers")) if top.get("layers") else None
-    aggregation_name = (top.get("aggregation") or "continuation_token").lower()
-    store_device = top.get("device") or top.get("store_device") or "cpu"
-    dtype = top.get("dtype")
-    save_dir = Path(top["save_dir"]) if top.get("save_dir") else None
-    return_full_sequence = to_bool(top.get("return_full_sequence", "false")) if "return_full_sequence" in top else False
-    normalize_layers = to_bool(top.get("normalize_layers", "false")) if "normalize_layers" in top else False
-    interactive = to_bool(top.get("interactive", "false")) if "interactive" in top else False
-    plan_only = to_bool(top.get("plan-only", top.get("plan_only", "false"))) if ( "plan-only" in top or "plan_only" in top ) else False
-    confirm = to_bool(top.get("confirm", "true")) if "confirm" in top else True
-
-    # Convert kwargs
-    loader_kwargs = parse_kv([f"{k}={v}" for k, v in loader_kv_raw.items()])
-    method_kwargs = parse_kv([f"{k}={v}" for k, v in method_kv_raw.items()])
-
-    # Interactive wizard
-    if interactive:
-        if loaders_location:
-            DataLoaderRotator.discover_loaders(loaders_location)
-        if not loader:
-            options = [d["name"] for d in DataLoaderRotator.list_loaders()]
-            loader = typer.prompt("Choose data loader", default=(options[0] if options else "custom"))
-        if loader and loader.lower() == "custom":
-            echo(Panel(
-                "[b]Custom loader arguments[/]\n\n"
-                "• path (str)  [required]\n"
-                "• split_ratio (float | None)\n"
-                "• seed (int | None)\n"
-                "• training_limit (int | None)\n"
-                "• testing_limit (int | None)",
-                title="custom.load(...)",
-            ) if HAS_RICH else
-                None
-            )
-            if "path" not in loader_kwargs:
-                loader_kwargs["path"] = typer.prompt("Path to dataset JSON (required)")
-            for name, cast, default in [
-                ("split_ratio", float, ""),
-                ("seed", int, ""),
-                ("training_limit", int, ""),
-                ("testing_limit", int, ""),
-            ]:
-                if name not in loader_kwargs:
-                    val = typer.prompt(f"{name} (optional)", default=default)
-                    if str(val).strip() != "":
-                        try:
-                            loader_kwargs[name] = cast(val)
-                        except Exception:
-                            loader_kwargs[name] = val
-        if not method_name:
-            method_name = typer.prompt("Choose steering method (see list-methods)", default="caa")
-        if layers is None:
-            layers = parse_layers(typer.prompt("Layers (e.g., '10..12', '5,7,9' or leave empty for all)", default="") or None)
-        if "aggregation" not in top:
-            aggregation_name = typer.prompt("Aggregation (see list-aggregations)", default="continuation_token")
-        if "dtype" not in top:
-            dtype = typer.prompt("Activation dtype (float32/float16/bfloat16 or blank)", default="") or None
-        if "device" not in top and "store_device" not in top:
-            store_device = typer.prompt("Device to store activations on (cpu / cuda / cuda:0 / ...)", default="cpu")
-        if "normalize_layers" not in top:
-            normalize_layers = typer.confirm("Normalize activations per layer?", default=True)
-        if "return_full_sequence" not in top:
-            return_full_sequence = typer.confirm("Return full [T,H] sequence per layer?", default=False)
-        if "save_dir" not in top:
-            default_out = os.path.abspath("./steering_output")
-            p = typer.prompt("Save directory for artifacts (blank to skip saving)", default=default_out)
-            if p.strip():
-                save_dir = Path(p)
-        if "plan-only" not in top and "plan_only" not in top:
-            plan_only = typer.confirm("Only show the plan and code preview?", default=False)
-        if "confirm" not in top:
-            confirm = typer.confirm("Confirm before running?", default=True)
-
-    # Validate dtype
-    if dtype not in DTYPE_MAP:
-        raise typer.BadParameter("dtype must be one of: float32, float16, bfloat16")
-
-    # Validate aggregation
-    try:
-        agg = aggs.pick(aggregation_name)
-    except ValueError as ex:
-        raise typer.BadParameter(str(ex)) from ex
-
-    # Plan
-    _show_plan(
-        model=model,
-        loader=loader,
-        loaders_location=loaders_location,
-        loader_kwargs=loader_kwargs,
-        method_name=method_name,
-        method_kwargs=method_kwargs,
-        layers=layers,
-        aggregation_name=aggregation_name,
-        store_device=store_device,
-        dtype=dtype,
-        return_full_sequence=return_full_sequence,
-        normalize_layers=normalize_layers,
-        save_dir=save_dir,
-    )
-
-    if plan_only:
-        return
-
-    if confirm and not typer.confirm("Proceed with training?", default=True):
-        typer.echo("Aborted.")
-        raise typer.Exit(code=1)
-
-    # -- Model -----------------------------------------------------------------
-    typer.echo(f"[+] Loading model: {model}")
-    from wisent.core.models.wisent_model import WisentModel  # type: ignore
-    wmodel = WisentModel(model_name=model, layers={}, device=store_device)
-
-    # -- Data loader -----------------------------------------------------------
-    from wisent.cli.data_loaders.data_loader_rotator import DataLoaderRotator  # type: ignore
-    if loaders_location:
-        DataLoaderRotator.discover_loaders(loaders_location)
-    dl_rot = DataLoaderRotator(loader=loader, loaders_location=loaders_location or "wisent_guard.core.data_loaders.loaders")
-    typer.echo(f"[+] Using data loader: {loader or '(default)'}")
-    load_result = dl_rot.load(**loader_kwargs)
-    pair_set = load_result["train_qa_pairs"]
-    typer.echo(f"[+] Loaded training pairs: {len(pair_set)} (task_type={load_result['task_type']})")
-
-    # -- Steering method -------------------------------------------------------
-    method_inst = _resolve_method(method_name, methods_location)
-    name_shown = getattr(method_inst, "name", type(method_inst).__name__)
-    typer.echo(f"[+] Steering method: {name_shown}")
-
-    # -- Trainer ---------------------------------------------------------------
-    from wisent.core.trainers.steering_trainer import WisentSteeringTrainer  # type: ignore
-    torch_dtype = None if dtype is None else __import__("torch").__dict__[DTYPE_MAP[dtype]]
-    trainer = WisentSteeringTrainer(
-        model=wmodel,
-        pair_set=pair_set,
-        steering_method=method_inst,
-        store_device=store_device,
-        dtype=torch_dtype,
-    )
-
-    result = trainer.run(
-        layers_spec=layers,
-        method_kwargs=method_kwargs,
-        aggregation=agg,
-        return_full_sequence=return_full_sequence,
-        normalize_layers=normalize_layers,
-        save_dir=save_dir,
-    )
-
-    typer.echo("\n=== Training Summary ===")
-    typer.echo(json.dumps(result.metadata, indent=2))
-    if save_dir is not None:
-        typer.echo(f"\nArtifacts saved in: {Path(save_dir).resolve()}\n")

wisent/cli/wisent_cli/main.py

@@ -1,93 +0,0 @@
-from __future__ import annotations
-from typing import Optional
-
-import typer
-
-from wisent.cli.wisent_cli.version import APP_NAME, APP_VERSION
-from wisent.cli.wisent_cli.ui import print_banner
-from wisent.cli.wisent_cli.commands.listing import app as listing_app
-from wisent.cli.wisent_cli.commands.train_cmd import app as train_app
-from wisent.cli.wisent_cli.commands.help_cmd import app as help_router_app
-from wisent.cli.wisent_cli.shell import app as shell_app
-
-app = typer.Typer(
-    no_args_is_help=True,
-    add_completion=False,
-    rich_markup_mode="markdown",
-    help=(
-        "[bold]Wisent Guard[/] – steerable activations / steering vectors.\n"
-        "Collect activations, train steering vectors, and inspect loaders & methods.\n\n"
-        "Natural commands (no dashes) + `help <topic>` and a `wisent` shell."
-    ),
-)
-
-app.add_typer(listing_app, name="list")
-app.add_typer(train_app,   name="")           # attach commands directly (e.g., train)
-app.add_typer(help_router_app, name="")       # help router lives at root (help ...)
-app.add_typer(shell_app,  name="shell")
-
-STATE = {"verbose": False}
-
-@app.callback(invoke_without_command=True)
-def _main_callback(
-    ctx: typer.Context,
-    version: Optional[bool] = typer.Option(None, "--version", "-V", help="Show version and exit."),
-    no_banner: bool = typer.Option(False, "--no-banner", help="Disable the startup banner."),
-    verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose logging."),
-    logo_width: Optional[int] = typer.Option(None, "--logo-width", help="Width of the wisent badge (28–96)."),
-):
-    """
-    Welcome to **Wisent Guard**.
-
-    Examples:
-      • `wisent help train`
-      • `wisent train model meta-llama/Llama-3.2-1B-Instruct loader custom path ./custom.json training_limit 5 method caa`
-      • `wisent list list-methods`      (methods)
-      • `wisent list list-loaders`      (loaders)
-      • `wisent list list-aggregations` (aggregations)
-      • `wisent shell start`            (interactive 'wisent' prompt)
-    """
-    if version:
-        typer.echo(f"{APP_NAME} {APP_VERSION}")
-        raise typer.Exit()
-
-    STATE["verbose"] = verbose
-    if not no_banner and (ctx.invoked_subcommand is None or ctx.info_name in {"--help", None}):
-        print_banner(APP_NAME, width=logo_width or 48)
-
-    if ctx.invoked_subcommand is None:
-        typer.echo(ctx.get_help())
-        raise typer.Exit()
-
-@app.command("instructions")
-def instructions():
-    msg = """
-[bold]Quickstart (no-dash style)[/]
-
-• Preview a run without executing:
-  [b]wisent train model gpt2 plan-only true[/]
-
-• Discover components:
-  [b]wisent list list-methods[/] | [b]wisent list list-loaders[/] | [b]wisent list list-aggregations[/]
-
-• Get help:
-  [b]wisent help train[/], [b]wisent help method caa[/], [b]wisent help loader custom[/]
-
-• Full training example (natural):
-  [b]wisent train model meta-llama/Llama-3.2-1B-Instruct loader custom path ./wisent_guard/cli/cli_examples/custom_dataset.json training_limit 5 \\
-        method caa layers 10..12 aggregation continuation_token device cuda dtype float16 save_dir ./steering_output normalize_layers true[/]
-
-Tip: add [b]interactive true[/] for a guided wizard.
-"""
-    try:
-        from rich.panel import Panel
-        from wisent.cli.wisent_cli.ui import echo
-        echo(Panel.fit(msg, title="Instructions", border_style="green"))
-    except Exception:
-        print(msg)
-
-def run():
-    app()
-
-if __name__ == "__main__":
-    run()

wisent/cli/wisent_cli/shell.py

@@ -1,80 +0,0 @@
-from __future__ import annotations
-import shlex
-import typer
-
-from wisent.cli.wisent_cli.ui import print_banner, echo
-
-try:
-    from rich.panel import Panel
-    HAS_RICH = True
-except Exception:
-    HAS_RICH = False
-
-__all__ = ["app", "start"]
-
-app = typer.Typer(help="Interactive shell")
-
-def _run_cli_line(line: str) -> None:
-    from typer.main import get_command
-    from wisent.cli.wisent_cli.main import app as root_app
-    click_cmd = get_command(root_app)
-    args = shlex.split(line)
-    try:
-        click_cmd.main(args=args, standalone_mode=False, prog_name="wisent")
-    except SystemExit as e:
-        if e.code not in (0, None):
-            raise
-
-@app.command("start")
-def start(
-    logo_width: int = typer.Option(48, "--logo-width", "-w", help="Logo width for the banner (28–96)."),
-    show_banner: bool = typer.Option(True, "--banner/--no-banner", help="Show banner when the shell starts."),
-):
-    """
-    Launch the interactive **wisent** shell.
-
-    Inside the shell, run:
-      • `help` / `help train` / `help method caa`
-      • `instructions`
-      • `train model ...` (no dashes)
-
-    Type `exit`, `quit`, or press `Ctrl-D` to leave.
-    """
-    if show_banner:
-        print_banner("Wisent Guard", logo_width)
-
-    hint = "Type 'help', 'help train', 'instructions', or any command like 'train model ...'. Type 'exit' to quit."
-    if HAS_RICH:
-        echo(Panel(hint, title="Welcome to the wisent shell", border_style="green"))
-    else:
-        print(hint)
-
-    GREEN, OFF = ("\x1b[32m", "\x1b[0m")
-    while True:
-        try:
-            if HAS_RICH:
-                from rich.console import Console
-                line = Console().input("[bold green]wisent[/] ")
-            else:
-                line = input(f"{GREEN}wisent{OFF} ")
-        except (EOFError, KeyboardInterrupt):
-            typer.echo("\nBye.")
-            break
-
-        line = line.strip()
-        if not line:
-            continue
-        if line in {"exit", "quit", "q"}:
-            typer.echo("Bye.")
-            break
-        if line in {"help", "--help", "-h"}:
-            _run_cli_line("--help")
-            continue
-        if line.startswith("help "):
-            _run_cli_line(line)
-            continue
-        if line in {"instructions", "--instructions"}:
-            _run_cli_line("instructions")
-            continue
-
-        _run_cli_line(line)

wisent/cli/wisent_cli/ui.py

@@ -1,69 +0,0 @@
-from __future__ import annotations
-
-from typing import Any
-
-try:
-    from rich.console import Console
-    from rich.panel import Panel
-    from rich.text import Text
-    HAS_RICH = True
-    _console = Console()
-except Exception:  # pragma: no cover
-    HAS_RICH = False
-    _console = None
-
-__all__ = ["print_banner", "echo"]
-
-def _render_wisent_mark(width: int = 48) -> str:
-    logo = """
-.................  .:--++*##%%%%##**+=-:.  .................
-..             .:=*%@@@@@@@%%%%%%%@@@@@@%*=:.             ..
-.           .-*%@@@%#+=-::.........:-=+#%@@@%*=.           .
-.         -*%@@@#=:.                    .:=*%@@@*-.        .
-.      .-#@@@*=.                            .-*@@@#-.      .
-.     :#@@@*:                                  :+%@@#-     .
-.   .+@@@*:                                      :+@@@+.   .
-.  .*@@@@%*=:.                                     -%@@#:  .
-. .#@@#=*%@@@%*-:.                                  .#@@%: .
-..*@@%.  .-+#@@@@#+-:.                               .*@@%..
-.=@@@-       :-+#@@@@%*=:.                            .%@@*.
-:#@@+           .:-+#@@@@%#+=:.                        -@@@-
-=@@@:                .-=*%@@@@%#+=:..                  .#@@+
-+@@@*=:.                 .:-+*%@@@@%#*=-:..             *@@+
-+@@@@@@#+-..                  .:-=*#@@@@@%#*+--..       +@@+
-+@@#-+%@@@%:                        .:-=*#%@@@@@%#*+=-:.*@@+
-=@@%. .=@@@:                             ..:-=+#%%@@@@@%@@@+
-:%@@=  :@@@-                                    ..::-=+#@@@=
-.+@@%. .#@@*                                           +@@#:
-..%@@*. =@@@:                                         =@@@-.
-. :%@@*..#@@#.                         .:..          =@@@= .
-.  :%@@*.:%@@*.                       :#@@%#*+=-::..+@@@=  .
-.   :#@@%-:%@@#:                    .+@@@#%%@@@@@@%%@@%-   .
-.    .+@@@*=#@@%-                 .=%@@%=...::-=#@@@@*.    .
-.      :*@@@#%@@@*:             .=%@@@+.     .:*%@@#-      .
-.        :+%@@@@@@@*-.       :=*@@@%+.    .-+%@@@*-.       .
-.          .=*%@@@@@@#+:.:-+#@@@%*-. .:-+#%@@@#+:          .
-.             .-+#%@@@@@@@@@@@@#*+**#@@@@@%*=:.            .
-..............   ..-=+*#%%%@@@@@@@@%%#*=-:.   ..............
- ...................  ....:::::::::.... ................... 
-""".strip("\n")
-    return "\n".join(line.center(width) for line in logo.splitlines())
-
-
-def print_banner(title: str, width: int | None = None) -> None:
-    if HAS_RICH:
-        usable = 64 if width is None else width
-        badge = _render_wisent_mark(usable)
-        _console.print(Panel.fit(Text(badge, style="green"), title=f"[b green]{title}[/]", border_style="green"))
-        _console.print("[dim]Steering vectors & activation tooling[/]\n")
-    else:
-        GREEN, OFF = "\x1b[32m", "\x1b[0m"
-        print(GREEN + _render_wisent_mark(width or 48) + OFF)
-        print(f"{title} – Steering vectors & activation tooling\n")
-
-
-def echo(obj: Any) -> None:
-    if HAS_RICH:
-        _console.print(obj)
-    else:
-        print(obj)

wisent/cli/wisent_cli/util/aggregations.py

@@ -1,43 +0,0 @@
-from __future__ import annotations
-from typing import Dict, Optional
-
-__all__ = ["agg_cls", "descriptions", "name_map", "pick"]
-
-def agg_cls():
-    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
-    return ActivationAggregationStrategy
-
-def descriptions() -> Dict[object, str]:
-    A = agg_cls()
-    return {
-        A.CHOICE_TOKEN: "Target A/B choice tokens (multiple choice).",
-        A.CONTINUATION_TOKEN: "Use the first token of the continuation.",
-        A.LAST_TOKEN: "Always select the last token.",
-        A.FIRST_TOKEN: "Always select the first token.",
-        A.MEAN_POOLING: "Aggregate by mean over all tokens.",
-        A.MAX_POOLING: "Aggregate by max over all tokens.",
-    }
-
-def name_map() -> Dict[str, object]:
-    A = agg_cls()
-    d = descriptions()
-    mapping = {s.name.lower(): s for s in d}
-    mapping.update({
-        "cont": A.CONTINUATION_TOKEN,
-        "choice": A.CHOICE_TOKEN,
-        "mean": A.MEAN_POOLING,
-        "max": A.MAX_POOLING,
-        "first": A.FIRST_TOKEN,
-        "last": A.LAST_TOKEN,
-    })
-    return mapping
-
-def pick(name: Optional[str]):
-    if not name:
-        return agg_cls().CONTINUATION_TOKEN
-    key = name.strip().lower()
-    mapping = name_map()
-    if key not in mapping:
-        valid = ", ".join(sorted(mapping.keys()))
-        raise ValueError(f"Unknown aggregation {name!r}. Valid: {valid}")
-    return mapping[key]