You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
from imaginairy.utils.model_manager import load_tensors
|
|
|
|
|
|
def dotstrings_to_nested_dictionaries(list_of_dotstrings):
|
|
"""given a list of dotstrings, return a nested dictionary."""
|
|
nested_dict = {}
|
|
for dotstring in list_of_dotstrings:
|
|
keys = dotstring.split(".")
|
|
d = nested_dict
|
|
for key in keys[:-1]:
|
|
if key not in d:
|
|
d[key] = {}
|
|
d = d[key]
|
|
d[keys[-1]] = None
|
|
return nested_dict
|
|
|
|
|
|
def display_nested_dictionary(nested_dictionary, num_levels_deep=2):
|
|
"""given a nested dictionary, print it out."""
|
|
|
|
def _display_nested_dictionary(d, level=0):
|
|
if level >= num_levels_deep:
|
|
return
|
|
for k, v in d.items():
|
|
print(" " * level + k)
|
|
if isinstance(v, dict):
|
|
_display_nested_dictionary(v, level + 1)
|
|
|
|
_display_nested_dictionary(nested_dictionary)
|
|
|
|
|
|
def display_weights_structure(weights_path, num_levels_deep=2):
|
|
"""given a weights path, display the structure of the weights."""
|
|
print(
|
|
f"Displaying weights structure for {weights_path} to {num_levels_deep} levels deep"
|
|
)
|
|
data = load_tensors(weights_path)
|
|
display_nested_dictionary(
|
|
dotstrings_to_nested_dictionaries(data.keys()), num_levels_deep=num_levels_deep
|
|
)
|