zhifu gao
2025-04-22 2c2fb5e1eb1185a081e3507c2aa5c3aafaa2bb6d
funasr/register.py
@@ -3,9 +3,11 @@
from dataclasses import dataclass
import re
@dataclass
class RegisterTables:
    """Registry system for classes."""
    model_classes = {}
    frontend_classes = {}
    specaug_classes = {}
@@ -46,16 +48,19 @@
    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.info(f"New registry table added: {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:
                raise ValueError(f"Key {registry_key} already exists in {register_tables_key}")
                logging.debug(
                    f"Key {registry_key} already exists in {register_tables_key}, re-register"
                )
            registry[registry_key] = target_class
@@ -78,4 +83,5 @@
        return decorator
tables = RegisterTables()