Coverage for /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/equine/load_equine_model.py: 100%

1# Copyright 2024, MASSACHUSETTS INSTITUTE OF TECHNOLOGY 

2# Subject to FAR 52.227-11 – Patent Rights – Ownership by the Contractor (May 2014). 

3# SPDX-License-Identifier: MIT 

4 

5from .equine import Equine 

6from .equine_gp import EquineGP 

7from .equine_protonet import EquineProtonet 

8from .utils import load_checkpoint 

9 

10 

11def load_equine_model(model_path: str) -> Equine: 

12 """ 

13 Attempt to load an EQUINE model from a file 

14 

15 Parameters 

16 ---------- 

17 model_path : str 

18 The path to the model file 

19 

20 Returns 

21 ------- 

22 Equine 

23 The loaded EQUINE model 

24 

25 Raises 

26 ------ 

27 ValueError 

28 If the model type is unknown, or if the file is not in the current 

29 EQUINE format (files from earlier releases must first be rewritten with 

30 ``python -m equine.upgrade``). 

31 """ 

32 model_save = load_checkpoint(model_path) 

33 model_type = model_save["train_summary"]["modelType"] 

34 

35 if model_type == "EquineProtonet": 

36 model = EquineProtonet._from_checkpoint(model_save) 

37 elif model_type == "EquineGP": 

38 model = EquineGP._from_checkpoint(model_save) 

39 else: 

40 raise ValueError(f"Unknown model type '{model_type}'") 

41 return model