Coverage for mindsdb / integrations / handlers / flaml_handler / flaml_handler.py: 0%

1import dill 

2import pandas as pd 

3from mindsdb.integrations.libs.base import BaseMLEngine 

4from typing import Dict, Optional 

5from type_infer.api import infer_types 

6from flaml import AutoML 

7 

8 

9class FLAMLHandler(BaseMLEngine): 

10 name = "FLAML" 

11 

12 def create(self, target: str, df: Optional[pd.DataFrame] = None, args: Optional[Dict] = None) -> None: 

13 if args is None: 

14 args = {} 

15 

16 if df is not None: 

17 target_dtype = infer_types(df, 0).to_dict()["dtypes"][target] 

18 model = AutoML(verbose=0) 

19 

20 if target_dtype in ['binary', 'categorical', 'tags']: 

21 model.fit(X_train=df.drop(columns=[target]), 

22 y_train=df[target], 

23 task='classification', 

24 **args.get('using')) 

25 

26 elif target_dtype in ['integer', 'float', 'quantity']: 

27 model.fit(X_train=df.drop(columns=[target]), 

28 y_train=df[target], 

29 task='regression', 

30 **args.get('using')) 

31 

32 self.model_storage.json_set('args', args) 

33 self.model_storage.file_set('model', dill.dumps(model)) 

34 

35 else: 

36 raise Exception( 

37 "Data is empty!!" 

38 ) 

39 

40 def predict(self, df: pd.DataFrame, args: Optional[Dict] = None) -> pd.DataFrame: 

41 

42 model = dill.loads(self.model_storage.file_get("model")) 

43 target = self.model_storage.json_get('args').get("target") 

44 

45 results = pd.DataFrame(model.predict(df), columns=[target]) 

46 

47 return results