Add *args, **kwargs to read data method in reader interface (#10)

Author: mpozniak95

dataset_reader/ann_h5_multi_reader.py

@@ -52,7 +52,7 @@ def read_queries(self) -> Iterator[Query]:
                 )
 
     def read_data(
-        self, start_idx: int = 0, end_idx: int = None, chunk_size: int = 10_000
+        self, start_idx: int = 0, end_idx: int = None, chunk_size: int = 10_000, *args, **kwargs
     ) -> Iterator[Record]:
         """
         Reads the 'train' data vectors from multiple HDF5 files based on the specified range.

dataset_reader/ann_h5_reader.py

@@ -27,7 +27,7 @@ def read_queries(self) -> Iterator[Query]:
                 expected_scores=expected_scores.tolist(),
             )
 
-    def read_data(self) -> Iterator[Record]:
+    def read_data(self, *args, **kwargs) -> Iterator[Record]:
         data = h5py.File(self.path)
 
         for idx, vector in enumerate(data["train"]):

dataset_reader/base_reader.py

@@ -18,7 +18,7 @@ class Query:
 
 
 class BaseReader:
-    def read_data(self) -> Iterator[Record]:
+    def read_data(self, *args, **kwargs) -> Iterator[Record]:
         raise NotImplementedError()
 
     def read_queries(self) -> Iterator[Query]:

dataset_reader/json_reader.py

@@ -60,7 +60,7 @@ def read_queries(self) -> Iterator[Query]:
 
             yield Query(vector=vector, meta_conditions=None, expected_result=neighbours)
 
-    def read_data(self) -> Iterator[Record]:
+    def read_data(self, *args, **kwargs) -> Iterator[Record]:
         for idx, (vector, payload) in enumerate(
             zip(self.read_vectors(), self.read_payloads())
         ):