Coverage for ml_workbench/model.py: 25%

1from __future__ import annotations 

2 

3from collections.abc import Mapping 

4from dataclasses import dataclass 

5from importlib import import_module 

6from typing import TYPE_CHECKING, Any 

7 

8if TYPE_CHECKING: 8 ↛ 9line 8 didn't jump to line 9 because the condition on line 8 was never true

9 from .config import YamlConfig 

10 

11 

12@dataclass(frozen=True) 

13class ModelSpec: 

14 name: str 

15 type: str 

16 description: str | None 

17 params: dict[str, Any] 

18 tuning: dict[str, Any] 

19 

20 

21class Model: 

22 """Encapsulates a model definition from the configuration. 

23 

24 Expected YAML structure: 

25 

26 models: 

27 my_model: 

28 description: "..." # optional 

29 type: "package.module.Class" # required 

30 params: { ... } # optional, free-form dict 

31 tuning: { ... } # optional, free-form dict 

32 

33 Notes 

34 ----- 

35 - The ``params`` and ``tuning`` sections are model-specific and are stored as given. 

36 - Use ``instantiate()`` to create the model class instance with provided params. 

37 """ 

38 

39 def __init__(self, name: str, config: YamlConfig) -> None: 

40 self.name = name 

41 self.config = config 

42 

43 models_section = self._get_models_section(config) 

44 if name not in models_section: 

45 raise KeyError(f"Model '{name}' not found in configuration") # noqa: TRY003 

46 

47 raw = models_section[name] 

48 if not isinstance(raw, Mapping): 

49 raise TypeError(f"Model '{name}' specification must be a mapping") # noqa: TRY003 

50 

51 model_type = raw.get("type") 

52 if not isinstance(model_type, str) or not model_type: 

53 raise ValueError(f"Model '{name}' is missing required 'type' string field") # noqa: TRY003 

54 

55 description = raw.get("description") 

56 params = raw.get("params", {}) 

57 tuning = raw.get("tuning", {}) 

58 

59 if not isinstance(params, Mapping): 

60 raise TypeError(f"Model '{name}'.params must be a mapping if provided") # noqa: TRY003 

61 if not isinstance(tuning, Mapping): 

62 raise TypeError(f"Model '{name}'.tuning must be a mapping if provided") # noqa: TRY003 

63 

64 self.spec = ModelSpec( 

65 name=name, 

66 type=model_type, 

67 description=description if isinstance(description, str) else None, 

68 params=dict(params), 

69 tuning=dict(tuning), 

70 ) 

71 

72 @staticmethod 

73 def _get_models_section(config: YamlConfig) -> Mapping[str, Any]: 

74 models = config.get_data().get("models") 

75 if not isinstance(models, Mapping): 

76 raise KeyError("No 'models' section found in configuration") # noqa: TRY003 

77 return models 

78 

79 @staticmethod 

80 def _import_string(path: str) -> tuple[object, str]: 

81 """Import a dotted path like 'package.module.Class'. 

82 

83 Returns the imported module object and the class name. 

84 """ 

85 if "." not in path: 

86 raise ValueError( # noqa: TRY003 

87 f"Invalid model type path '{path}'. Expected 'package.module.Class'" 

88 ) 

89 module_path, class_name = path.rsplit(".", 1) 

90 module = import_module(module_path) 

91 return module, class_name 

92 

93 def instantiate(self, **overrides: Any) -> Any: 

94 """Instantiate the model class using configured params. 

95 

96 Parameters 

97 ---------- 

98 overrides: Any 

99 Optional overrides for parameters when constructing the model. 

100 """ 

101 module, cls_name = self._import_string(self.spec.type) 

102 cls = getattr(module, cls_name) 

103 init_kwargs = {**self.spec.params, **overrides} 

104 return cls(**init_kwargs) 

105 

106 def to_dict(self) -> dict[str, Any]: 

107 return { 

108 "name": self.spec.name, 

109 "type": self.spec.type, 

110 "description": self.spec.description, 

111 "params": dict(self.spec.params), 

112 "tuning": dict(self.spec.tuning), 

113 } 

114 

115 @classmethod 

116 def list_model_names(cls, config: YamlConfig) -> list[str]: 

117 models = cls._get_models_section(config) 

118 return list(models.keys()) 

119 

120 @classmethod 

121 def verify_config(cls, config: YamlConfig) -> None: 

122 """Validate that all models are well-formed. 

123 

124 Checks 

125 ------ 

126 - models section exists and is a mapping 

127 - each model has a non-empty string 'type' 

128 - params and tuning, if present, are mappings 

129 """ 

130 models = config.get_data().get("models") 

131 if models is None: 

132 return 

133 if not isinstance(models, Mapping): 

134 raise TypeError("'models' section must be a mapping") # noqa: TRY003 

135 

136 for model_name, raw in models.items(): 

137 if not isinstance(raw, Mapping): 

138 raise TypeError( # noqa: TRY003 

139 f"Model '{model_name}' specification must be a mapping" 

140 ) 

141 model_type = raw.get("type") 

142 if not isinstance(model_type, str) or not model_type: 

143 raise ValueError( # noqa: TRY003 

144 f"Model '{model_name}' is missing required 'type' string field" 

145 ) 

146 params = raw.get("params", {}) 

147 tuning = raw.get("tuning", {}) 

148 if not isinstance(params, Mapping): 

149 raise TypeError( # noqa: TRY003 

150 f"Model '{model_name}'.params must be a mapping if provided" 

151 ) 

152 if not isinstance(tuning, Mapping): 

153 raise TypeError( # noqa: TRY003 

154 f"Model '{model_name}'.tuning must be a mapping if provided" 

155 ) 

156 

157 

158__all__ = ["Model", "ModelSpec"]