Coverage for src/sensai/util/multiprocessing.py: 44%

1import pandas as pd 

2 

3from ..vector_model import VectorModel 

4 

5 

6class VectorModelWithSeparateFeatureGeneration: 

7 def __init__(self, vector_model: VectorModel): 

8 self.vectorModel = vector_model 

9 self.featureGen = vector_model.get_feature_generator() 

10 self.vectorModel.set_feature_generator(None) 

11 

12 def __str__(self): 

13 return self.vectorModel.__str__() 

14 

15 class IntermediateFittingStep: 

16 def __init__(self, vector_model: VectorModel, x: pd.DataFrame, y: pd.DataFrame): 

17 self.y = y 

18 self.x = x 

19 self.vector_model = vector_model 

20 

21 def execute(self) -> VectorModel: 

22 self.vector_model.fit(self.x, self.y) 

23 return self.vector_model 

24 

25 def __str__(self): 

26 return f"{self.__class__.__name__} for {self.vector_model}" 

27 

28 class PredictFinaliser: 

29 def __init__(self, vector_model: VectorModel, x: pd.DataFrame): 

30 self.X = x 

31 self.vectorModel = vector_model 

32 

33 def execute(self) -> pd.DataFrame: 

34 return self.vectorModel.predict(self.X) 

35 

36 def __str__(self): 

37 return f"{self.__class__.__name__} for {self.vectorModel}" 

38 

39 def fit_start(self, x, y) -> 'VectorModelWithSeparateFeatureGeneration.IntermediateFittingStep': 

40 x = self.featureGen.fit_generate(x, y) 

41 return self.IntermediateFittingStep(self.vectorModel, x, y) 

42 

43 def predict_start(self, x: pd.DataFrame): 

44 x = self.featureGen.generate(x) 

45 return self.PredictFinaliser(self.vectorModel, x) 

46 

47 def fit_end(self, vector_model) -> VectorModel: 

48 vector_model._featureGenerator = self.featureGen 

49 return vector_model