Add TokenCountBatchSampler from COMBO

combo/data/samplers/__init__.py

-from .base_sampler import Sampler
+from .batch_sampler import BatchSampler
 from .samplers import TokenCountBatchSampler

combo/data/samplers/base_sampler.py

-class Sampler:
-    pass

combo/data/samplers/batch_sampler.py

+"""
+Re-implemented AllenNLP BatchSampler
+https://github.com/allenai/allennlp/blob/main/allennlp/data/samplers/batch_sampler.py
+"""
+from typing import Sequence, Iterable, List, Optional
+from torch import Tensor
+
+
+class BatchSampler:
+    def get_batch_indices(self, instances: Sequence[Tensor]) -> Iterable[List[int]]:
+        raise NotImplementedError
+
+    def get_num_batches(self, instances: Sequence[Tensor]) -> int:
+        raise NotImplementedError
+
+    def get_batch_size(self) -> Optional[int]:
+        """
+        Not all `BatchSamplers` define a consistent `batch_size`, but those that
+        do should override this method.
+        """
+        return None

combo/data/samplers/samplers.py

+"""
+Adapted from COMBO
+Author: Mateusz Klimaszewski
+"""
+
 from typing import List
 
 import numpy as np
 
-from combo.data.samplers import Sampler
+from combo.data.samplers import BatchSampler
 
 
-class TokenCountBatchSampler(Sampler):
+class TokenCountBatchSampler(BatchSampler):
 
     def __init__(self, dataset, word_batch_size: int = 2500, shuffle_dataset: bool = True):
         self._index = 0