Fixed multi-module typed signature optimizer

dspy/functional/functional.py

@@ -82,6 +82,9 @@ def __init__(self, signature, max_retries=3):
     def copy(self) -> "TypedPredictor":
         return TypedPredictor(self.signature, self.max_retries)
 
+    def __repr__(self):
+        return f"TypedPredictor({self.signature})"
+
     @staticmethod
     def _make_example(type_) -> str:
         # Note: DSPy will cache this call so we only pay the first time TypedPredictor is called.

dspy/primitives/module.py

@@ -1,4 +1,5 @@
 import copy
+from collections import deque
 from collections.abc import Generator
 
 import ujson
@@ -9,45 +10,48 @@ def __init__(self):
         pass
 
     def named_parameters(self):
-        """
-        Unlike PyTorch, handles (non-recursive) lists of parameters too.
-        """
-
+        """Unlike PyTorch, handles lists of parameters too."""
         from dspy.predict.parameter import Parameter
 
-        visited = set()
-        named_parameters = []
-
-        def add_parameter(param_name, param_value):
-            if isinstance(param_value, Parameter) and id(param_value) not in visited:
-                visited.add(id(param_value))
-                named_parameters.append((param_name, param_value))
-
-        for name, value in self.__dict__.items():
-            if isinstance(value, Parameter):
-                add_parameter(name, value)
+        # Remove the 'self.' prefix from the names
+        return [(name[5:], param) for name, param in self.named_sub_modules(Parameter)]
 
-            elif isinstance(value, BaseModule):
-                # When a sub-module is pre-compiled, keep it frozen.
-                if not getattr(value, "_compiled", False):
-                    for sub_name, param in value.named_parameters():
-                        add_parameter(f"{name}.{sub_name}", param)
+    def named_sub_modules(self, type_=None, skip_compiled=False) -> Generator[tuple[str, "BaseModule"], None, None]:
+        """Find all sub-modules in the module, as well as their names.
 
-            elif isinstance(value, (list, tuple)):
-                for idx, item in enumerate(value):
-                    add_parameter(f"{name}[{idx}]", item)
-
-            elif isinstance(value, dict):
-                for key, item in value.items():
-                    add_parameter(f"{name}['{key}']", item)
-
-        return named_parameters
-
-    def named_sub_modules(self, root_name="base") -> Generator[tuple[str, "BaseModule"], None, None]:
-        yield root_name, self
-        for name, value in self.__dict__.items():
-            if isinstance(value, BaseModule):
-                yield from value.named_sub_modules(root_name=f"{root_name}.{name}")
+        Say self.children[4]['key'].sub_module is a sub-module. Then the name will be
+        'children[4][key].sub_module'. But if the sub-module is accessible at different
+        paths, only one of the paths will be returned.
+        """
+        if type_ is None:
+            type_ = BaseModule
+
+        queue = deque([("self", self)])
+        seen = {id(self)}
+
+        def add_to_queue(name, item):
+            if id(item) not in seen:
+                seen.add(id(item))
+                queue.append((name, item))
+
+        while queue:
+            name, item = queue.popleft()
+            if isinstance(item, type_):
+                yield name, item
+
+            if isinstance(item, BaseModule):
+                if skip_compiled and getattr(item, "_compiled", False):
+                    continue
+                for sub_name, sub_item in item.__dict__.items():
+                    add_to_queue(f"{name}.{sub_name}", sub_item)
+
+            elif isinstance(item, (list, tuple)):
+                for i, sub_item in enumerate(item):
+                    add_to_queue(f"{name}[{i}]", sub_item)
+
+            elif isinstance(item, dict):
+                for key, sub_item in item.items():
+                    add_to_queue(f"{name}[{key}]", sub_item)
 
     def parameters(self):
         return [param for _, param in self.named_parameters()]

dspy/teleprompt/signature_opt_typed.py

@@ -272,13 +272,15 @@ def optimize_signature(
         pass
     elif strategy == "best":
         i = scores.index(max(scores))
+        if verbose:
+            print(f"Best signature: {i} with score: {scores[i]}")
         for name, p in named_predictors:
             p.signature = candidates[name][i].to_signature()
     else:
         raise ValueError(f"Invalid strategy: {strategy}")
 
     return OptimizerResult(
         program=module,
-        signatures=[{name: sigs[i].to_signature()} for name, sigs in candidates.items() for i in range(n_iterations)],
+        signatures=[{name: sigs[i].to_signature() for name, sigs in candidates.items()} for i in range(n_iterations)],
         scores=scores,
     )