import logging
|
import inspect
|
from dataclasses import dataclass
|
import re
|
|
|
@dataclass
|
class RegisterTables:
|
"""Registry system for classes."""
|
|
model_classes = {}
|
frontend_classes = {}
|
specaug_classes = {}
|
normalize_classes = {}
|
encoder_classes = {}
|
decoder_classes = {}
|
joint_network_classes = {}
|
predictor_classes = {}
|
stride_conv_classes = {}
|
tokenizer_classes = {}
|
dataloader_classes = {}
|
batch_sampler_classes = {}
|
dataset_classes = {}
|
index_ds_classes = {}
|
|
def print(self, key: str = None) -> None:
|
"""Print registered classes."""
|
print("\ntables: \n")
|
fields = vars(self)
|
headers = ["register name", "class name", "class location"]
|
for classes_key, classes_dict in fields.items():
|
if classes_key.endswith("_meta") and (key is None or key in classes_key):
|
print(f"----------- ** {classes_key.replace('_meta', '')} ** --------------")
|
metas = []
|
for register_key, meta in classes_dict.items():
|
metas.append(meta)
|
metas.sort(key=lambda x: x[0])
|
data = [headers] + metas
|
col_widths = [max(len(str(item)) for item in col) for col in zip(*data)]
|
|
for row in data:
|
print(
|
"| "
|
+ " | ".join(str(item).ljust(width) for item, width in zip(row, col_widths))
|
+ " |"
|
)
|
print("\n")
|
|
def register(self, register_tables_key: str, key: str = None) -> callable:
|
"""Decorator to register a class."""
|
|
def decorator(target_class):
|
if not hasattr(self, register_tables_key):
|
setattr(self, register_tables_key, {})
|
logging.debug(f"New registry table added: {register_tables_key}")
|
|
registry = getattr(self, register_tables_key)
|
registry_key = key if key is not None else target_class.__name__
|
|
if registry_key in registry:
|
logging.debug(
|
f"Key {registry_key} already exists in {register_tables_key}, re-register"
|
)
|
|
registry[registry_key] = target_class
|
|
register_tables_key_meta = register_tables_key + "_meta"
|
if not hasattr(self, register_tables_key_meta):
|
setattr(self, register_tables_key_meta, {})
|
registry_meta = getattr(self, register_tables_key_meta)
|
|
class_file = inspect.getfile(target_class)
|
class_line = inspect.getsourcelines(target_class)[1]
|
pattern = r"^.+/funasr/"
|
class_file = re.sub(pattern, "funasr/", class_file)
|
meta_data = [
|
registry_key,
|
target_class.__name__,
|
f"{class_file}:{class_line}",
|
]
|
registry_meta[registry_key] = meta_data
|
return target_class
|
|
return decorator
|
|
|
tables = RegisterTables()
|
