refactor: refactor dataloaders

Author: boomb0om

DPF/dataloaders/dataloader_utils.py

@@ -1,12 +1,12 @@
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 
 # TODO(review) - логика работы непонятна совсем, для чего метод нужен, нужны пояснения + рефактор (выглядит как что-то ненужное)
 # default identical preprocessing function for FilesDataset and ShardsDataset
-def identical_preprocess_function(column2bytes: Dict[str, bytes], data: Dict[str, str]) -> Any:
- return column2bytes, data
+def identical_preprocess_function(modality2data: Dict[str, Union[bytes, Any]], data: Dict[str, str]) -> Any:
+ return modality2data, data
 
 
 # identical collate function for pytorch dataloader
-def identical_collate_fn(x):
+def identical_collate_fn(x: Any) -> Any:
  return x

DPF/dataloaders/files_dataset.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import pandas as pd
 from torch.utils.data import Dataset
@@ -8,7 +8,7 @@
 from DPF.filesystems.filesystem import FileSystem
 
 
-class FilesDataset(Dataset):
+class FilesDataset(Dataset[Tuple[bool, Any]]):
  """
  Dataset class to read "raw" files
  """
@@ -19,7 +19,7 @@ def __init__(
  df: pd.DataFrame,
  datatypes: List[Union[ShardedDataType, FileDataType, ColumnDataType]],
  meta_columns: Optional[List[str]] = None,
- preprocess_function: Callable[[Dict[str, bytes], Dict[str, str]], Any] = identical_preprocess_function,
+ preprocess_function: Callable[[Dict[str, Union[bytes, Any]], Dict[str, str]], Any] = identical_preprocess_function,
  # TODO(review) - на ошибке надо выбрасывать ошибку, а не возвращать None, и в дальнейшем эту ошибку обрабатывать прикладом, использующим этот класс
  return_none_on_error: bool = False
  ):
@@ -50,7 +50,7 @@ def __init__(
  self.column2modality = {}
  for d in self.datatypes:
  if isinstance(d, ColumnDataType):
- self.column2modality[d.modality.column] = d.modality.key
+ self.column2modality[d.column_name] = d.modality.key
  elif isinstance(d, (ShardedDataType, FileDataType)):
  self.path_column2modality[d.modality.path_column] = d.modality.key
  else:
@@ -64,10 +64,10 @@ def __init__(
  self.preprocess_f = preprocess_function
  self.return_none_on_error = return_none_on_error
 
- def __len__(self):
+ def __len__(self) -> int:
  return len(self.data_to_iterate)
 
- def __getitem__(self, idx):
+ def __getitem__(self, idx: int) -> Tuple[bool, Any]:
  data = {
  self.columns[c]: item for c, item in enumerate(self.data_to_iterate[idx])
  }

DPF/dataloaders/shards_dataset.py

@@ -1,7 +1,7 @@
 import itertools
 import os
 import tarfile
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 
 import pandas as pd
 import torch
@@ -12,7 +12,7 @@
 from DPF.filesystems.filesystem import FileSystem
 
 
-class ShardsDataset(IterableDataset):
+class ShardsDataset(IterableDataset[Tuple[bool, Any]]):
  """
  Dataset class for shards format (files in tar archives)
  """
@@ -24,7 +24,7 @@ def __init__(
  split2archive_path: Dict[str, str],
  datatypes: List[Union[ShardedDataType, ColumnDataType]],
  meta_columns: Optional[List[str]] = None,
- preprocess_function: Callable[[Dict[str, bytes], Dict[str, str]], Any] = identical_preprocess_function,
+ preprocess_function: Callable[[Dict[str, Union[bytes, Any]], Dict[str, str]], Any] = identical_preprocess_function,
  return_none_on_error: bool = False
  ):
  """
@@ -44,9 +44,9 @@ def __init__(
  Preprocessing function for data. First argument of the preprocess_f is mapping from modality name to bytes
  and the second argument is mapping from meta_column name to its value.
  return_none_on_error: bool = False
- Whether to return None if error during reading file occures
+ Whether to return None if error during reading file occurs
  """
- super(ShardsDataset).__init__()
+ super().__init__()
  self.filesystem = filesystem
 
  self.datatypes = datatypes
@@ -57,7 +57,7 @@ def __init__(
  self.column2modality = {}
  for d in self.datatypes:
  if isinstance(d, ColumnDataType):
- self.column2modality[d.modality.column] = d.modality.key
+ self.column2modality[d.column_name] = d.modality.key
  elif isinstance(d, ShardedDataType):
  self.path_column2modality[d.modality.path_column] = d.modality.key
  else:
@@ -76,10 +76,10 @@ def __init__(
  self.preprocess_f = preprocess_function
  self.return_none_on_error = return_none_on_error
 
- def __len__(self):
+ def __len__(self) -> int:
  return self.total_samples
 
- def __iter__(self):
+ def __iter__(self) -> Iterator[Tuple[bool, Any]]:
  worker_info = torch.utils.data.get_worker_info()
  worker_total_num = worker_info.num_workers if worker_info is not None else None
  worker_id = worker_info.id if worker_info is not None else None
@@ -100,12 +100,12 @@ def __iter__(self):
  filename = os.path.basename(data[col])
  if self.return_none_on_error:
  try:
- file_bytes = tar.extractfile(filename).read()
+ file_bytes = tar.extractfile(filename).read() # type: ignore
  except Exception:
  file_bytes = None
  is_ok = False
  else:
- file_bytes = tar.extractfile(filename).read()
+ file_bytes = tar.extractfile(filename).read() # type: ignore
  modality2data[modality] = file_bytes
  # read data from columns
  for col in self.column2modality.keys():