fennol.models.modules

  1from typing import Sequence, Tuple, Dict, Optional, Any
  2import flax.linen as nn
  3from inspect import isclass, ismodule
  4from pkgutil import iter_modules
  5import importlib
  6import importlib.util
  7import os
  8import glob
  9
 10### python modules where FENNIX Modules are defined ###
 11from . import misc,physics, embeddings,preprocessing
 12
 13
 14MODULES: Dict[str, nn.Module] = {}
 15PREPROCESSING: Dict = {}
 16
 17
 18def available_fennix_modules():
 19    return list(MODULES.keys())
 20
 21def available_fennix_preprocessing():
 22    return list(PREPROCESSING.keys())
 23
 24
 25def register_fennix_module(module: nn.Module, FID: Optional[str] = None):
 26    if FID is not None:
 27        names = [FID.upper()]
 28    else:
 29        if not hasattr(module, "FID"):
 30            print(f"Warning: module {module.__name__} does not have a FID field and no explicit FID was provided. Module was NOT registered.")
 31        if not (isinstance(
 32            module.FID, str
 33        ) or isinstance(module.FID,tuple)):
 34            print(f"Warning: module {module.__name__} has an invalid FID field. Module was NOT registered.")
 35        if isinstance(module.FID, str):
 36            names = [module.FID.upper()]
 37        else:
 38            names = [fid.upper() for fid in module.FID]
 39    for name in names:
 40        if name in MODULES and MODULES[name] != module:
 41            raise ValueError(
 42                f"A different module identified as '{name}' is already registered !"
 43            )
 44        MODULES[name] = module
 45
 46def register_fennix_preprocessing(module, FPID: Optional[str] = None):
 47    if FPID is not None:
 48        names = [FPID.upper()]
 49    else:
 50        if not hasattr(module, "FPID"):
 51            print(f"Warning: module {module.__name__} does not have a FPID field and no explicit FPID was provided. Module was NOT registered.")
 52        if not (isinstance(
 53            module.FPID, str
 54        ) or isinstance(module.FPID,tuple)):
 55            print(f"Warning: module {module.__name__} has an invalid FPID field. Module was NOT registered.")
 56        if isinstance(module.FPID, str):
 57            names = [module.FPID.upper()]
 58        else:
 59            names = [fid.upper() for fid in module.FPID]
 60    for name in names:
 61        if name in PREPROCESSING and PREPROCESSING[name] != module:
 62            raise ValueError(
 63                f"A different module identified as '{name}' is already registered !"
 64            )
 65        PREPROCESSING[name] = module
 66
 67
 68def register_fennix_modules(module, recurs=0, max_recurs=2):
 69    if ismodule(module) and hasattr(module,"__path__"):
 70        for _, name, _ in iter_modules(module.__path__):
 71            sub_module = __import__(f"{module.__name__}.{name}", fromlist=[""])
 72            register_fennix_modules(sub_module, recurs=recurs + 1, max_recurs=max_recurs)
 73    for m in module.__dict__.values():
 74        if isclass(m) and issubclass(m, nn.Module) and m != nn.Module:
 75            if hasattr(m, "FID"):
 76                register_fennix_module(m)
 77        elif isclass(m) and hasattr(m, "FPID"):
 78            register_fennix_preprocessing(m)
 79
 80
 81### REGISTER DEFAULT MODULES #####################
 82for mods in [misc, physics, embeddings,preprocessing]:
 83    register_fennix_modules(mods)
 84module_path = os.environ.get("FENNOL_MODULES_PATH","").split(":")
 85for path in module_path:
 86    if os.path.exists(path):
 87        for file in glob.glob(f"{path}/*.py"):
 88            spec = importlib.util.spec_from_file_location("custom_modules", file)
 89            module = importlib.util.module_from_spec(spec)
 90            spec.loader.exec_module(module)
 91            register_fennix_modules(module)
 92
 93##################################################
 94
 95def get_modules_documentation():
 96    doc = {}
 97    for name, module in MODULES.items():
 98        doc[name] = module.__doc__
 99    return doc
100
101
102class FENNIXModules(nn.Module):
103    r"""Sequential module that applies a sequence of FENNIX modules.
104
105    Attributes:
106        layers (Sequence[Tuple[nn.Module, Dict]]): Sequence of tuples (module, parameters) to apply.
107
108    """
109
110    layers: Sequence[Tuple[nn.Module, Dict]]
111
112    def __post_init__(self):
113        if not isinstance(self.layers, Sequence):
114            raise ValueError(
115                f"'layers' must be a sequence, got '{type(self.layers).__name__}'."
116            )
117        super().__post_init__()
118
119    @nn.compact
120    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
121        if not self.layers:
122            raise ValueError(f"Empty Sequential module {self.name}.")
123
124        outputs = inputs
125        for layer, prms in self.layers:
126            outputs = layer(**prms)(outputs)
127        return outputs