Load a model from an installed Python module (e.g. transformers).
(
module_name: str, class_name: str, pretrained: Optional[str] = None
)
| 247 | |
| 248 | |
| 249 | def _load_model_from_module( |
| 250 | module_name: str, class_name: str, pretrained: Optional[str] = None |
| 251 | ) -> nn.Module: |
| 252 | """Load a model from an installed Python module (e.g. transformers).""" |
| 253 | try: |
| 254 | module = importlib.import_module(module_name) |
| 255 | except ImportError as e: |
| 256 | raise ImportError( |
| 257 | f"Cannot import module '{module_name}'. Is it installed? Error: {e}" |
| 258 | ) from e |
| 259 | |
| 260 | if not hasattr(module, class_name): |
| 261 | raise AttributeError( |
| 262 | f"Class '{class_name}' not found in module '{module_name}'." |
| 263 | ) |
| 264 | |
| 265 | cls = getattr(module, class_name) |
| 266 | |
| 267 | if pretrained: |
| 268 | # HuggingFace-style: cls.from_pretrained(...) |
| 269 | if hasattr(cls, "from_pretrained"): |
| 270 | try: |
| 271 | model = cls.from_pretrained(pretrained, torch_dtype="auto") |
| 272 | except Exception as e: |
| 273 | raise RuntimeError( |
| 274 | f"Failed to load pretrained model '{pretrained}' via " |
| 275 | f"{class_name}.from_pretrained(): {e}" |
| 276 | ) from e |
| 277 | else: |
| 278 | raise RuntimeError( |
| 279 | f"{class_name} does not have a from_pretrained() method." |
| 280 | ) |
| 281 | else: |
| 282 | try: |
| 283 | model = cls() |
| 284 | except TypeError as e: |
| 285 | raise RuntimeError( |
| 286 | f"Could not instantiate {class_name}() with no arguments: {e}. " |
| 287 | "Use --pretrained for HuggingFace models." |
| 288 | ) from e |
| 289 | |
| 290 | return model |
| 291 | |
| 292 | |
| 293 | def load_model(args: argparse.Namespace) -> Tuple[nn.Module, str]: |