feature: implement method take for datastream

Author: samedii

datastream/datastream.py

@@ -65,6 +65,9 @@ def __init__(
             )
         )
 
+    def __len__(self):
+        return len(self.sampler)
+
     @staticmethod
     def merge(datastreams_and_ns: Tuple[Union[
         Datastream[T],
@@ -213,6 +216,16 @@ def sample_proportion(
             self.sampler.sample_proportion(proportion),
         )
 
+    def take(
+        self: Datastream[T],
+        n_samples: int,
+    ) ->  Datastream[T]:
+        '''
+        Like :func:`Datastream.sample_proportion` but specify the number of
+        samples instead of a proportion.
+        '''
+        return self.sample_proportion(min(1, n_samples / len(self)))
+
     def state_dict(self) -> Dict:
         '''Get state of datastream. Useful for checkpointing sample weights.'''
         return dict(sampler=self.sampler.state_dict())

datastream/samplers.py

@@ -24,7 +24,7 @@ def __init__(self, length, proportion=1.0, replacement=False):
             replacement=replacement,
             sampler=torch.utils.data.WeightedRandomSampler(
                 torch.ones(length).double(),
-                num_samples=int(length * proportion),
+                num_samples=max(int(length * proportion), 1),
                 replacement=replacement,
             )
         )