Spaces:
Running on Zero
Running on Zero
File size: 4,608 Bytes
fed6c68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | # Copyright 2025 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re
from typing import List
import torch.nn as nn
def set_module_from_path(model: nn.Module, path: str, value: any):
attrs = path.split(".")
if len(attrs) == 1:
setattr(model, attrs[0], value)
else:
next_obj = getattr(model, attrs[0])
set_module_from_path(next_obj, ".".join(attrs[1:]), value)
def get_module_from_path(model: nn.Module, path: str):
attrs = path.split(".")
if len(attrs) == 1:
return getattr(model, attrs[0])
else:
next_obj = getattr(model, attrs[0])
return get_module_from_path(next_obj, ".".join(attrs[1:]))
def is_same_module_from_path(model: nn.Module, module_name: str, path: str):
attrs = path.split(".")
if len(attrs) == 1:
return getattr(model, attrs[0]).__class__.__name__ == module_name
else:
next_obj = getattr(model, attrs[0])
return is_same_module_from_path(next_obj, module_name, ".".join(attrs[1:]))
def check_all_fqn_match(path_patterns: List[str], path_keys: List[str]):
"""
Check
"""
assert isinstance(path_patterns, list), f"path_patterns must be a list, got {type(path_patterns)}"
assert isinstance(path_keys, (list, tuple)), f"path_keys must be a list or tuple, got {type(path_keys)}"
if len(path_patterns) != len(path_keys):
return False
regex_list = []
for pattern in path_patterns:
regex_str = re.escape(pattern).replace(r"\*", r"(\d+)")
regex_str = f"^{regex_str}$"
regex_list.append((pattern, re.compile(regex_str)))
used_patterns = set()
expected_num = None # the first matched number
for key in path_keys:
matched = False
for p, regex in regex_list:
if p in used_patterns:
continue
match = regex.match(key)
if match:
current_num = match.group(1)
if expected_num is None:
expected_num = current_num
elif current_num != expected_num:
return False
used_patterns.add(p)
matched = True
break
if not matched:
return False
return True
def check_any_fqn_match(path_patterns: List[str], path_key: str, return_idx: bool = False, prefix: str = None):
assert isinstance(path_patterns, list), f"path_patterns must be a list, got {type(path_patterns)}"
assert isinstance(path_key, str), f"path_key must be a str, got {type(path_key)}"
if prefix:
path_patterns = [".".join([prefix, pattern]) for pattern in path_patterns]
regex_list = []
for pattern in path_patterns:
regex_str = re.escape(pattern).replace(r"\*", r"(\d+)")
regex_str = f"^{regex_str}$"
regex_list.append(re.compile(regex_str))
for idx, regex in enumerate(regex_list):
match = regex.match(path_key)
if match:
return idx if return_idx else True
return -1 if return_idx else False
def check_fqn_match(fqn_pattern: str, fqn: str, prefix: str = None):
assert isinstance(fqn_pattern, str), f"fqn_pattern must be a str, got {type(fqn_pattern)}"
assert isinstance(fqn, str), f"fqn must be a str, got {type(fqn)}"
if prefix:
fqn_pattern = [".".join([prefix, pattern]) for pattern in fqn_pattern]
regex_str = re.escape(fqn_pattern).replace(r"\*", r".*")
regex_str = f"^{regex_str}$"
regex = re.compile(regex_str)
match = regex.match(fqn)
return match
def sort_fqn_by_submodule_first(fqn_list: list[str]) -> list[str]:
"""
Sort FQN list purely by string nesting relationship (ignore depth calculation)
"""
def _fqn_nesting_compare(a: str, b: str) -> int:
if a in b and a != b:
return 1
elif b in a and b != a:
return -1
else:
return 0
from functools import cmp_to_key
sorted_list = sorted(fqn_list, key=cmp_to_key(_fqn_nesting_compare))
return sorted_list
|