diff --git a/lora_optimizer.py b/lora_optimizer.py index 3405ae7..6ac14d5 100644 --- a/lora_optimizer.py +++ b/lora_optimizer.py @@ -166,6 +166,93 @@ def _resolve_arch_preset(arch_override, detected_arch): return key, _ARCH_PRESETS[key] +def _parse_merge_formula(formula_str, n_loras): + """ + Parse a merge formula string into a tree structure. + + Syntax: + expr = term (('+') term)* + term = atom (':' weight)? + atom = NUMBER | '(' expr ')' + weight = FLOAT + + Numbers are 1-indexed LoRA positions. Returns a tree of: + {"type": "leaf", "index": int, "weight": float|None} + {"type": "group", "children": list, "weight": float|None} + + Raises ValueError on malformed input or out-of-range indices. + """ + formula_str = formula_str.strip() + if not formula_str: + raise ValueError("Empty merge formula") + + pos = [0] # mutable position cursor + + def _skip_ws(): + while pos[0] < len(formula_str) and formula_str[pos[0]] == ' ': + pos[0] += 1 + + def _parse_weight(): + _skip_ws() + if pos[0] < len(formula_str) and formula_str[pos[0]] == ':': + pos[0] += 1 # skip ':' + _skip_ws() + start = pos[0] + while pos[0] < len(formula_str) and (formula_str[pos[0]].isdigit() or formula_str[pos[0]] == '.'): + pos[0] += 1 + if pos[0] == start: + raise ValueError(f"Expected weight after ':' at position {pos[0]}") + return float(formula_str[start:pos[0]]) + return None + + def _parse_atom(): + _skip_ws() + if pos[0] >= len(formula_str): + raise ValueError("Unexpected end of formula") + + if formula_str[pos[0]] == '(': + pos[0] += 1 # skip '(' + node = _parse_expr() + _skip_ws() + if pos[0] >= len(formula_str) or formula_str[pos[0]] != ')': + raise ValueError(f"Expected ')' at position {pos[0]}") + pos[0] += 1 # skip ')' + weight = _parse_weight() + node["weight"] = weight + return node + + # Must be a number + start = pos[0] + while pos[0] < len(formula_str) and formula_str[pos[0]].isdigit(): + pos[0] += 1 + if pos[0] == start: + raise ValueError(f"Unexpected character '{formula_str[pos[0]]}' at position {pos[0]}") + index_1based = int(formula_str[start:pos[0]]) + if index_1based < 1 or index_1based > n_loras: + raise ValueError(f"LoRA index {index_1based} out of range (have {n_loras} LoRAs)") + weight = _parse_weight() + return {"type": "leaf", "index": index_1based - 1, "weight": weight} + + def _parse_expr(): + children = [_parse_atom()] + while True: + _skip_ws() + if pos[0] < len(formula_str) and formula_str[pos[0]] == '+': + pos[0] += 1 # skip '+' + children.append(_parse_atom()) + else: + break + if len(children) == 1: + return children[0] + return {"type": "group", "children": children, "weight": None} + + result = _parse_expr() + _skip_ws() + if pos[0] != len(formula_str): + raise ValueError(f"Unexpected content at position {pos[0]}: '{formula_str[pos[0]:]}'") + return result + + class LoRAStack: """ Node for creating a LoRA stack (input format for LoRAOptimizer). diff --git a/tests/test_lora_optimizer.py b/tests/test_lora_optimizer.py index cde4721..aa8c0fd 100644 --- a/tests/test_lora_optimizer.py +++ b/tests/test_lora_optimizer.py @@ -674,5 +674,70 @@ class LoRASettingsNodeTests(unittest.TestCase): "_DEFAULTS keys don't match INPUT_TYPES keys") + # --- Merge formula parser tests --- + + def test_parse_merge_formula_simple(self): + """Simple flat formula parses to group of leaves.""" + tree = lora_optimizer._parse_merge_formula("1 + 2 + 3", 3) + self.assertEqual(tree["type"], "group") + self.assertEqual(len(tree["children"]), 3) + for i, child in enumerate(tree["children"]): + self.assertEqual(child["type"], "leaf") + self.assertEqual(child["index"], i) + + def test_parse_merge_formula_nested(self): + """Nested formula parses to tree with sub-group.""" + tree = lora_optimizer._parse_merge_formula("(1+2) + 3", 3) + self.assertEqual(tree["type"], "group") + self.assertEqual(len(tree["children"]), 2) + sub = tree["children"][0] + self.assertEqual(sub["type"], "group") + self.assertEqual(len(sub["children"]), 2) + leaf3 = tree["children"][1] + self.assertEqual(leaf3["type"], "leaf") + self.assertEqual(leaf3["index"], 2) + + def test_parse_merge_formula_weights(self): + """Weights are parsed from :N.N suffix.""" + tree = lora_optimizer._parse_merge_formula("(1+2):0.6 + 3:0.4", 3) + self.assertAlmostEqual(tree["children"][0]["weight"], 0.6) + self.assertAlmostEqual(tree["children"][1]["weight"], 0.4) + + def test_parse_merge_formula_deep_nesting(self): + """Deep nesting: ((1+2)+3) + 4.""" + tree = lora_optimizer._parse_merge_formula("((1+2)+3) + 4", 4) + self.assertEqual(tree["type"], "group") + self.assertEqual(len(tree["children"]), 2) + inner = tree["children"][0] + self.assertEqual(inner["type"], "group") + self.assertEqual(len(inner["children"]), 2) + innermost = inner["children"][0] + self.assertEqual(innermost["type"], "group") + self.assertEqual(len(innermost["children"]), 2) + + def test_parse_merge_formula_single_item(self): + """Single item is valid.""" + tree = lora_optimizer._parse_merge_formula("1", 1) + self.assertEqual(tree["type"], "leaf") + self.assertEqual(tree["index"], 0) + + def test_parse_merge_formula_out_of_range(self): + """Out of range index raises ValueError.""" + with self.assertRaises(ValueError): + lora_optimizer._parse_merge_formula("1 + 5", 3) + + def test_parse_merge_formula_malformed(self): + """Malformed formula raises ValueError.""" + with self.assertRaises(ValueError): + lora_optimizer._parse_merge_formula("((1+2", 3) + + def test_parse_merge_formula_empty(self): + """Empty/whitespace formula raises ValueError.""" + with self.assertRaises(ValueError): + lora_optimizer._parse_merge_formula("", 3) + with self.assertRaises(ValueError): + lora_optimizer._parse_merge_formula(" ", 3) + + if __name__ == "__main__": unittest.main()