diff --git a/src/soup_cli/utils/adapter_arithmetic.py b/src/soup_cli/utils/adapter_arithmetic.py new file mode 100644 index 0000000..e9f8519 --- /dev/null +++ b/src/soup_cli/utils/adapter_arithmetic.py @@ -0,0 +1,248 @@ +"""LoRA adapter task-vector arithmetic (v0.71.34). + +Task arithmetic (arXiv:2212.04089) applied to LoRA adapters: add / scale / +NEGATE task vectors via an expression such as ``"coder + 0.5*math - toxic"``. + +The engine is **signed, un-normalized, element-wise** over the intersection of +``lora_A`` / ``lora_B`` tensor names (mirrors PEFT ``combination_type="linear"``): +``out[k] = Σ cᵢ·tensorᵢ[k]``. Same-rank inputs only — a shape mismatch on a shared +tensor is a rank mismatch and is rejected loudly (harmonize rank first). Exact +concatenation+SVD arithmetic for mixed-rank adapters is a future enhancement. + +No top-level torch/transformers/peft — the parser + numpy merge stay light. +""" + +from __future__ import annotations + +import json +import math +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence, Set, Tuple + +_MAX_EXPR_LEN = 4096 +_MAX_TERMS = 64 +_MAX_ADAPTER_CONFIG_BYTES = 256 * 1024 + +_NAME_RE = re.compile(r"[A-Za-z0-9_.\-]+") +_FLOAT_RE = re.compile(r"[0-9]+(?:\.[0-9]+)?(?:[eE][-+]?[0-9]+)?") + + +@dataclass(frozen=True) +class TaskTerm: + """A single ``coeff * adapter`` term of a task-arithmetic expression.""" + + name: str + coeff: float + + +@dataclass(frozen=True) +class ArithmeticReport: + """Result of a ``soup adapters arithmetic`` run.""" + + expression: str + terms: Tuple[TaskTerm, ...] + output_dir: str + merged_layers: int + skipped_layers: Tuple[str, ...] + base_model: str | None + + +def parse_expression(expr: str, known_names: Set[str]) -> list[TaskTerm]: + """Parse a task-arithmetic expression into signed ``TaskTerm``s. + + Grammar (NO ``eval``): ``expr := term (('+'|'-') term)*`` where + ``term := [sign] [coeff '*'] name`` (also ``name ['*' coeff]``). Signs fold + into the coefficient; an omitted coefficient is ``1.0``. Duplicate adapter + names sum their coefficients; a term summing to ``0.0`` is dropped. + + Raises: + ValueError: empty / over-length expression, unknown token, adapter name + not in ``known_names``, non-finite coefficient, all terms cancel, + or more than ``_MAX_TERMS`` distinct adapters. + """ + if not isinstance(expr, str): + raise TypeError("expression must be a string") + s = expr.strip() + if not s: + raise ValueError("empty expression") + if len(s) > _MAX_EXPR_LEN: + raise ValueError(f"expression too long (> {_MAX_EXPR_LEN} chars)") + + n = len(s) + i = 0 + coeffs: dict[str, float] = {} + order: list[str] = [] + seen_term = False + + def _skip_ws() -> None: + nonlocal i + while i < n and s[i] in " \t": + i += 1 + + _skip_ws() + while i < n: + _skip_ws() + if i >= n: + break + # Leading sign(s) — only valid before the first term or between terms. + sign = 1.0 + if s[i] in "+-": + while i < n and s[i] in "+- \t": + if s[i] == "-": + sign = -sign + i += 1 + elif seen_term: + # Two terms with no operator between them → malformed. + raise ValueError( + f"expected '+' or '-' between terms at pos {i}: {s[i:i + 8]!r}" + ) + _skip_ws() + if i >= n: + raise ValueError("expression ends with a dangling operator") + + coeff = 1.0 + name: str | None = None + m = _FLOAT_RE.match(s, i) + if m: + coeff = float(m.group()) + i = m.end() + _skip_ws() + if i < n and s[i] == "*": + i += 1 + _skip_ws() + nm = _NAME_RE.match(s, i) + if not nm: + raise ValueError( + f"expected adapter name after coefficient at pos {i}" + ) + name = nm.group() + i = nm.end() + else: + nm = _NAME_RE.match(s, i) + if not nm: + raise ValueError( + f"unexpected token at pos {i}: {s[i:i + 8]!r}" + ) + name = nm.group() + i = nm.end() + _skip_ws() + if i < n and s[i] == "*": + i += 1 + _skip_ws() + cm = _FLOAT_RE.match(s, i) + if not cm: + raise ValueError( + f"expected coefficient after '*' at pos {i}" + ) + coeff = float(cm.group()) + i = cm.end() + + if not math.isfinite(coeff): + raise ValueError("coefficient must be finite") + if name not in known_names: + raise ValueError( + f"unknown adapter name {name!r} " + f"(declare it with --adapter {name}=)" + ) + signed = sign * coeff + if name not in coeffs: + coeffs[name] = 0.0 + order.append(name) + if len(order) > _MAX_TERMS: + raise ValueError(f"too many distinct terms (> {_MAX_TERMS})") + coeffs[name] += signed + seen_term = True + _skip_ws() + + if not seen_term: + raise ValueError("expression has no terms") + terms = [TaskTerm(nm, coeffs[nm]) for nm in order if coeffs[nm] != 0.0] + if not terms: + raise ValueError("all terms cancelled to zero — nothing to merge") + return terms + + +def merge_task_arithmetic( + weights_list: Sequence[Mapping[str, Any]], + coeffs: Sequence[float], +) -> Tuple[dict[str, Any], Tuple[str, ...]]: + """Signed, un-normalized element-wise ``out[k] = Σ cᵢ·tensorᵢ[k]``. + + Operates over the intersection of tensor names; names present in only some + adapters are reported in ``skipped``. A shape mismatch on a *shared* name is + a rank mismatch and raises (same-rank contract). + """ + import numpy as np + + if len(weights_list) != len(coeffs): + raise ValueError( + f"weights_list ({len(weights_list)}) and coeffs " + f"({len(coeffs)}) length mismatch" + ) + if not weights_list: + raise ValueError("need at least one adapter") + + shared = set(weights_list[0].keys()) + all_keys = set(weights_list[0].keys()) + for w in weights_list[1:]: + shared &= set(w.keys()) + all_keys |= set(w.keys()) + + merged: dict[str, Any] = {} + for name in sorted(shared): + tensors = [np.asarray(w[name], dtype=np.float64) for w in weights_list] + if len({t.shape for t in tensors}) > 1: + raise ValueError( + f"rank/shape mismatch on {name!r} across adapters — task " + f"arithmetic requires same-rank adapters (harmonize the LoRA " + f"rank first, or merge with `soup adapters merge --strategy svd`)" + ) + acc = np.zeros_like(tensors[0]) + for c, t in zip(coeffs, tensors): + acc += float(c) * t + merged[name] = acc.astype(np.float32) + + skipped = tuple(sorted(all_keys - shared)) + return merged, skipped + + +def read_adapter_base(adapter_dir: str) -> str | None: + """Read ``base_model_name_or_path`` from an adapter's ``adapter_config.json``. + + Returns ``None`` when the config is absent or the field is missing. The read + is symlink-rejecting (O_NOFOLLOW) and size-capped at 256 KiB (mirrors + ``adapter_merge.write_merged_adapter``'s config-read guards). + """ + cfg_path = Path(adapter_dir) / "adapter_config.json" + if not os.path.lexists(str(cfg_path)): + return None + try: + fd = os.open(str(cfg_path), os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + except OSError as exc: + raise ValueError( + f"adapter_config.json unreadable: {type(exc).__name__}" + ) from exc + fh = None + try: + if os.fstat(fd).st_size > _MAX_ADAPTER_CONFIG_BYTES: + raise ValueError( + f"adapter_config.json exceeds {_MAX_ADAPTER_CONFIG_BYTES} byte cap" + ) + fh = os.fdopen(fd, "r", encoding="utf-8") + raw = fh.read() + finally: + if fh is not None: + fh.close() + else: + os.close(fd) + try: + data = json.loads(raw) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ValueError(f"adapter_config.json is not valid JSON: {exc}") from exc + if not isinstance(data, dict): + return None + base = data.get("base_model_name_or_path") + return base if isinstance(base, str) else None diff --git a/tests/test_v07134.py b/tests/test_v07134.py new file mode 100644 index 0000000..fdf42ed --- /dev/null +++ b/tests/test_v07134.py @@ -0,0 +1,181 @@ +"""v0.71.34 — Adapter algebra (task arithmetic) + LISA (#267). + +Covers: +* ``utils/adapter_arithmetic.py`` — expression parser + signed element-wise + task-vector merge + adapter base reader (no top-level torch). +* ``commands/adapters.py::arithmetic`` — ``soup adapters arithmetic``. +* ``config/schema.py`` — LISA fields + ``_validate_lisa_compat``. +* ``utils/lisa.py`` — ``LisaPolicy`` + ``LisaCallback`` (duck-typed). +* ``utils/peft_wiring.py::attach_lisa_callback`` + SFT trainer wiring. +""" + +from __future__ import annotations + +import ast +import json +from pathlib import Path + +import numpy as np +import pytest + + +# --------------------------------------------------------------------------- +# Task A1 — expression parser +# --------------------------------------------------------------------------- +class TestParseExpression: + def _names(self): + return {"coder", "math", "toxic"} + + def test_happy_add_scale_sub(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + terms = parse_expression("coder + 0.5*math - toxic", self._names()) + got = {t.name: t.coeff for t in terms} + assert got == {"coder": 1.0, "math": 0.5, "toxic": -1.0} + + def test_name_star_coeff(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + terms = parse_expression("coder*2", self._names()) + assert terms[0].name == "coder" and terms[0].coeff == 2.0 + + def test_leading_negative(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + terms = parse_expression("-coder + math", self._names()) + got = {t.name: t.coeff for t in terms} + assert got == {"coder": -1.0, "math": 1.0} + + def test_single_term_scale(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + terms = parse_expression("2*coder", self._names()) + assert len(terms) == 1 and terms[0].coeff == 2.0 + + def test_duplicate_names_sum(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + terms = parse_expression("coder + coder", self._names()) + assert len(terms) == 1 and terms[0].coeff == 2.0 + + def test_all_cancel_rejected(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + with pytest.raises(ValueError, match="cancel"): + parse_expression("coder - coder", self._names()) + + def test_empty_rejected(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + with pytest.raises(ValueError, match="empty"): + parse_expression(" ", self._names()) + + def test_unknown_name_rejected(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + with pytest.raises(ValueError, match="ghost"): + parse_expression("coder + ghost", self._names()) + + def test_injection_rejected(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + for bad in ['__import__("os")', "coder; rm -rf", "coder && ls", "coder | cat"]: + with pytest.raises(ValueError): + parse_expression(bad, self._names()) + + def test_over_length_rejected(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + with pytest.raises(ValueError, match="too long"): + parse_expression("coder+" * 5000 + "coder", self._names()) + + def test_non_finite_coeff_rejected(self): + from soup_cli.utils.adapter_arithmetic import parse_expression + + # "nan"/"inf" are names by charset, not floats — so they parse as + # unknown adapter names, not as coefficients. The finite guard defends + # against a hypothetical float token; assert the injection path rejects. + with pytest.raises(ValueError): + parse_expression("nan*coder", self._names()) + + def test_no_top_level_torch(self): + import soup_cli.utils.adapter_arithmetic as mod + + src = Path(mod.__file__).read_text(encoding="utf-8") + tree = ast.parse(src) + for node in tree.body: + if isinstance(node, (ast.Import, ast.ImportFrom)): + names = [] + if isinstance(node, ast.Import): + names = [a.name for a in node.names] + else: + names = [node.module or ""] + for nm in names: + assert nm.split(".")[0] not in { + "torch", + "transformers", + "peft", + }, f"top-level heavy import: {nm}" + + +# --------------------------------------------------------------------------- +# Task A2 — signed merge + base reader +# --------------------------------------------------------------------------- +class TestMergeTaskArithmetic: + def test_subtract(self): + from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic + + a = {"lora_A": np.ones((2, 3), dtype=np.float32)} + b = {"lora_A": np.full((2, 3), 4.0, dtype=np.float32)} + merged, skipped = merge_task_arithmetic([a, b], [1.0, -1.0]) + assert np.allclose(merged["lora_A"], -3.0) + assert skipped == () + + def test_scale(self): + from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic + + a = {"w": np.ones((2, 2), dtype=np.float32)} + merged, _ = merge_task_arithmetic([a], [2.5]) + assert np.allclose(merged["w"], 2.5) + + def test_mixed_rank_rejected(self): + from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic + + a = {"w": np.ones((2, 3), dtype=np.float32)} + b = {"w": np.ones((4, 3), dtype=np.float32)} + with pytest.raises(ValueError, match="rank"): + merge_task_arithmetic([a, b], [1.0, 1.0]) + + def test_disjoint_keys_skipped(self): + from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic + + a = {"shared": np.ones((2, 2), dtype=np.float32), "only_a": np.ones((1, 1))} + b = {"shared": np.ones((2, 2), dtype=np.float32), "only_b": np.ones((1, 1))} + merged, skipped = merge_task_arithmetic([a, b], [1.0, 1.0]) + assert "shared" in merged + assert set(skipped) == {"only_a", "only_b"} + + def test_length_mismatch_rejected(self): + from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic + + with pytest.raises(ValueError, match="length"): + merge_task_arithmetic([{"w": np.ones((1, 1))}], [1.0, 2.0]) + + +class TestReadAdapterBase: + def test_reads_base(self, tmp_path): + from soup_cli.utils.adapter_arithmetic import read_adapter_base + + d = tmp_path / "ad" + d.mkdir() + (d / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "meta/x"}), encoding="utf-8" + ) + assert read_adapter_base(str(d)) == "meta/x" + + def test_missing_returns_none(self, tmp_path): + from soup_cli.utils.adapter_arithmetic import read_adapter_base + + d = tmp_path / "ad" + d.mkdir() + assert read_adapter_base(str(d)) is None