refactor: refactor base filter classes

Author: boomb0om

DPF/filters/column_filter.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Dict, Any
 
 import numpy as np
 import pandas as pd
@@ -28,11 +28,11 @@ def schema(self) -> List[str]:
  pass
 
  @abstractmethod
- def process(self, row: dict) -> tuple:
+ def process(self, row: Dict[str, Any]) -> List[Any]:
  pass
 
- def __call__(self, df: pd.DataFrame) -> np.ndarray:
+ def __call__(self, df: pd.DataFrame) -> List[List[Any]]:
  pandarallel.initialize(nb_workers=self.workers)
- res = np.array(list(df[self.columns_to_process].parallel_apply(self.process, axis=1)))
+ res = list(df[self.columns_to_process].parallel_apply(self.process, axis=1))
  return res
 

DPF/filters/data_filter.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Tuple
 
 import pandas as pd
 from torch.utils.data import DataLoader, Dataset
@@ -51,32 +51,39 @@ def metadata_columns(self) -> List[str]:
  pass
 
  @abstractmethod
- def preprocess(self, modality2data: ModalityToDataMapping, metadata: Dict[str, str]):
+ def preprocess_data(
+ self,
+ modality2data: ModalityToDataMapping,
+ metadata: Dict[str, str]
+ ) -> Any:
  pass
 
  @abstractmethod
- def process_batch(self, batch) -> dict:
+ def process_batch(self, batch: List[Any]) -> Dict[str, List[Any]]:
  pass
 
  @staticmethod
- def _add_values_from_batch(main_dict: dict, batch_dict: dict):
+ def _add_values_from_batch(
+ main_dict: Dict[str, List[Any]],
+ batch_dict: Dict[str, List[Any]]
+ ) -> None:
  for k, v in batch_dict.items():
  main_dict[k].extend(v)
 
- def _generate_dict_from_schema(self):
+ def _get_dict_from_schema(self) -> Dict[str, List[Any]]:
  return {i: [] for i in self.schema}
 
- def run(self, dataset: Dataset) -> pd.DataFrame:
+ def run(self, dataset: Dataset[Tuple[bool, Any]]) -> pd.DataFrame:
  dataloader = DataLoader(dataset, collate_fn=identical_collate_fn, **self.dataloader_kwargs)
- df_labels = self._generate_dict_from_schema()
+ filter_results = self._get_dict_from_schema()
 
  for batch in tqdm(dataloader, disable=not self.pbar, position=self.pbar_position):
  # drop Nans
  batch_filtered = [b[1] for b in batch if b[0]]
  if len(batch_filtered) == 0:
  continue
 
- df_batch_labels = self.process_batch(batch_filtered)
- self._add_values_from_batch(df_labels, df_batch_labels)
+ filter_results_batch = self.process_batch(batch_filtered)
+ self._add_values_from_batch(filter_results, filter_results_batch)
 
- return pd.DataFrame(df_labels)
+ return pd.DataFrame(filter_results)

DPF/filters/multigpu_filter.py

@@ -1,5 +1,5 @@
 from multiprocessing import Manager, Process
-from typing import List, Type, Union
+from typing import List, Type, Union, Dict, Any
 
 import numpy as np
 import pandas as pd
@@ -21,13 +21,13 @@ def run_one_process(
  index: pd.Series,
  results: List[pd.DataFrame],
  filter_class: Type[DataFilter],
- filter_kwargs: dict,
- device: str,
- filter_run_kwargs: dict
-):
+ filter_kwargs: Dict[str, Any],
+ device: Union[str, torch.device],
+ filter_run_kwargs: Dict[str, Any]
+) -> None:
  reader = DatasetReader(filesystem=fs)
  processor = reader.from_df(config, df)
- datafilter = filter_class(**filter_kwargs, _pbar_position=i, device=device)
+ datafilter = filter_class(**filter_kwargs, _pbar_position=i, device=device) # type: ignore
  processor.apply_data_filter(datafilter, **filter_run_kwargs)
  res = processor.df
  res.set_index(index, inplace=True)
@@ -41,16 +41,22 @@ class MultiGPUDataFilter:
 
  def __init__(
  self,
- devices: List[Union[torch.device | str]],
- filter_class: type,
- filter_params: dict
+ devices: List[Union[torch.device, str]],
+ datafilter_class: Type[DataFilter],
+ datafilter_params: Dict[str, Any]
  ):
- self.filter_class = filter_class
- self.filter_params = filter_params
+ self.filter_class = datafilter_class
+ self.filter_params = datafilter_params
  self.devices = devices
  self.num_parts = len(devices)
 
- def run(self, df: pd.DataFrame, config: DatasetConfig, fs: FileSystem, filter_run_kwargs: dict) -> pd.DataFrame:
+ def run(
+ self,
+ df: pd.DataFrame,
+ config: DatasetConfig,
+ fs: FileSystem,
+ filter_run_kwargs: Dict[str, Any]
+ ) -> pd.DataFrame:
  manager = Manager()
  shared_results = manager.list()
 
@@ -63,7 +69,7 @@ def run(self, df: pd.DataFrame, config: DatasetConfig, fs: FileSystem, filter_ru
  fs,
  df_splits[i],
  i,
- df_splits[i].index,
+ df_splits[i].index, # type: ignore
  shared_results,
  self.filter_class,
  self.filter_params,