|
|
|
|
|
from ...optimizers.engine.registry import ParamRegistry |
|
|
from typing import Any, Optional, List, Dict |
|
|
|
|
|
|
|
|
class MiproRegistry(ParamRegistry): |
|
|
""" |
|
|
Extended ParamRegistry that supports storing input_names and output_names |
|
|
for each optimizable field. Compatible with all original track() usages. |
|
|
""" |
|
|
|
|
|
def track( |
|
|
self, |
|
|
root_or_obj: Any, |
|
|
path_or_attr: str = None, |
|
|
*, |
|
|
name: Optional[str] = None, |
|
|
input_names: Optional[List[str]] = None, |
|
|
output_names: Optional[List[str]] = None, |
|
|
input_descs: Optional[Dict[str, str]] = None, |
|
|
output_descs: Optional[Dict[str, str]] = None, |
|
|
): |
|
|
|
|
|
if isinstance(root_or_obj, (list, tuple)): |
|
|
for item in root_or_obj: |
|
|
if isinstance(item, dict): |
|
|
self.track(**item) |
|
|
elif isinstance(item, (list, tuple)): |
|
|
if len(item) == 7: |
|
|
self.track( |
|
|
item[0], item[1], |
|
|
name=item[2], |
|
|
input_names=item[3], |
|
|
output_names=item[4], |
|
|
input_descs=item[5], |
|
|
output_descs=item[6] |
|
|
) |
|
|
else: |
|
|
raise ValueError("Each tuple must be (obj, attr, name, input_names, output_names, input_descs, output_descs)") |
|
|
return self |
|
|
|
|
|
|
|
|
super().track(root_or_obj, path_or_attr, name=name) |
|
|
|
|
|
|
|
|
key = name or path_or_attr |
|
|
field = self.fields[key] |
|
|
field.input_names = input_names or [] |
|
|
field.output_names = output_names or [] |
|
|
field.input_descs = input_descs or {} |
|
|
field.output_descs = output_descs or {} |
|
|
|
|
|
return self |
|
|
|
|
|
def get_input_names(self, name: str) -> List[str]: |
|
|
"""Return the input_names for a registered field, or an empty list if not set.""" |
|
|
return getattr(self.fields[name], "input_names", None) or [] |
|
|
|
|
|
def get_output_names(self, name: str) -> List[str]: |
|
|
"""Return the output_names for a registered field, or an empty list if not set.""" |
|
|
return getattr(self.fields[name], "output_names", None) or [] |
|
|
|
|
|
def get_input_desc_dict(self, name: str) -> Dict[str, str]: |
|
|
"""Return the input_descs for a registered field, or an empty dict if not set.""" |
|
|
return getattr(self.fields[name], "input_descs", {}) |
|
|
|
|
|
def get_output_desc_dict(self, name: str) -> Dict[str, str]: |
|
|
"""Return the output_descs for a registered field, or an empty dict if not set.""" |
|
|
return getattr(self.fields[name], "output_descs", {}) |
|
|
|
|
|
def get_input_desc(self, name: str, input_name: str) -> str: |
|
|
"""Return the input_desc for a registered field, or an empty string if not set.""" |
|
|
return self.get_input_desc_dict(name).get(input_name, "") |
|
|
|
|
|
def get_output_desc(self, name: str, output_name: str) -> str: |
|
|
"""Return the output_desc for a registered field, or an empty string if not set.""" |
|
|
return self.get_output_desc_dict(name).get(output_name, "") |
|
|
|
|
|
def describe(self) -> Dict[str, Dict[str, Any]]: |
|
|
""" |
|
|
Returns a dict of all fields and their metadata, including input/output names if present. |
|
|
""" |
|
|
result = {} |
|
|
for name, field in self.fields.items(): |
|
|
result[name] = { |
|
|
"value": field.get(), |
|
|
"input_names": getattr(field, "input_names", None), |
|
|
"output_names": getattr(field, "output_names", None), |
|
|
"input_descs": getattr(field, "input_descs", {}), |
|
|
"output_descs": getattr(field, "output_descs", {}), |
|
|
} |
|
|
return result |
