Coverage for mindsdb / integrations / libs / ml_handler_process / predict_process.py: 30%

1import importlib 

2 

3from pandas import DataFrame 

4 

5import mindsdb.interfaces.storage.db as db 

6from mindsdb.interfaces.storage.model_fs import ModelStorage, HandlerStorage 

7from mindsdb.integrations.libs.ml_handler_process.handlers_cacher import handlers_cacher 

8from mindsdb.utilities.functions import mark_process 

9 

10 

11@mark_process(name='learn') 

12def predict_process(integration_id: int, predictor_record: db.Predictor, args: dict, 

13 module_path: str, ml_engine_name: str, dataframe: DataFrame) -> DataFrame: 

14 module = importlib.import_module(module_path) 

15 

16 if predictor_record.id not in handlers_cacher: 

17 handlerStorage = HandlerStorage(integration_id) 

18 modelStorage = ModelStorage(predictor_record.id) 

19 ml_handler = module.Handler( 

20 engine_storage=handlerStorage, 

21 model_storage=modelStorage, 

22 ) 

23 handlers_cacher[predictor_record.id] = ml_handler 

24 else: 

25 ml_handler = handlers_cacher[predictor_record.id] 

26 

27 if ml_engine_name == 'lightwood': 

28 args['code'] = predictor_record.code 

29 args['target'] = predictor_record.to_predict[0] 

30 args['dtype_dict'] = predictor_record.dtype_dict 

31 args['learn_args'] = predictor_record.learn_args 

32 

33 predictions = ml_handler.predict(dataframe, args) 

34 ml_handler.close() 

35 return predictions