mirror of https://github.com/razor-ai/soup.git
feat(adapters): task-vector arithmetic parser + signed merge (v0.71.34)
This commit is contained in:
Alpamys
parent
4e88f70691
commit
662ecef3a7
|
|
@ -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}=<path>)"
|
||||
)
|
||||
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
|
||||
|
|
@ -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
|
||||
Loading…
Reference in New Issue