Load a model from an installed Python module (e.g. 'transformers').
(module_name: str, class_name: str,
pretrained: Optional[str] = None, **kwargs)
| 127 | |
| 128 | |
| 129 | def load_model_from_module(module_name: str, class_name: str, |
| 130 | pretrained: Optional[str] = None, **kwargs) -> nn.Module: |
| 131 | """Load a model from an installed Python module (e.g. 'transformers').""" |
| 132 | try: |
| 133 | mod = importlib.import_module(module_name) |
| 134 | except ImportError as e: |
| 135 | raise ImportError( |
| 136 | f"Cannot import module '{module_name}'. Is it installed? Error: {e}" |
| 137 | ) |
| 138 | |
| 139 | if not hasattr(mod, class_name): |
| 140 | raise AttributeError( |
| 141 | f"Class '{class_name}' not found in module '{module_name}'." |
| 142 | ) |
| 143 | |
| 144 | cls = getattr(mod, class_name) |
| 145 | |
| 146 | if pretrained: |
| 147 | # HuggingFace-style: cls.from_pretrained(...) |
| 148 | if hasattr(cls, "from_pretrained"): |
| 149 | model = cls.from_pretrained(pretrained, **kwargs) |
| 150 | else: |
| 151 | raise AttributeError( |
| 152 | f"'{class_name}' has no 'from_pretrained' method. " |
| 153 | f"Cannot load pretrained weights from '{pretrained}'." |
| 154 | ) |
| 155 | else: |
| 156 | model = cls(**kwargs) |
| 157 | |
| 158 | return model |
| 159 | |
| 160 | |
| 161 | def load_model(args) -> nn.Module: |