Repository: adbar/simplemma

Issue:
Add protocol classes for classes that can be extended by the user
At the moment these are:
* DictionaryFactory
* Tokenizer
* TokenSampler

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi) Report
> Merging [#84](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi) (233ee6d) into [main](https://app.codecov.io/gh/adbar/simplemma/commit/e8decd8d7df5d799c73a6a472713642d49edcf3f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi) (e8decd8) will **increase** coverage by `0.03%`.
> The diff coverage is `97.91%`.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi).

```diff
@@            Coverage Diff             @@
##             main      #84      +/-   ##
==========================================
+ Coverage   96.47%   96.50%   +0.03%     
==========================================
  Files          17       30      +13     
  Lines         510      629     +119     
==========================================
+ Hits          492      607     +115     
- Misses         18       22       +4     
```


| [Impacted Files](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi) | Coverage Δ | |
|---|---|---|
| [simplemma/strategies/defaultrules/de.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3N0cmF0ZWdpZXMvZGVmYXVsdHJ1bGVzL2RlLnB5) | `100.00% <ø> (ø)` | |
| [simplemma/strategies/defaultrules/en.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3N0cmF0ZWdpZXMvZGVmYXVsdHJ1bGVzL2VuLnB5) | `92.30% <ø> (ø)` | |
| [simplemma/strategies/defaultrules/fi.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3N0cmF0ZWdpZXMvZGVmYXVsdHJ1bGVzL2ZpLnB5) | `87.50% <ø> (ø)` | |
| [simplemma/strategies/defaultrules/generic.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3N0cmF0ZWdpZXMvZGVmYXVsdHJ1bGVzL2dlbmVyaWMucHk=) | `100.00% <ø> (ø)` | |
| [simplemma/strategies/defaultrules/nl.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3N0cmF0ZWdpZXMvZGVmYXVsdHJ1bGVzL25sLnB5) | `91.66% <ø> (ø)` | |
| [simplemma/strategies/defaultrules/pl.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3N0cmF0ZWdpZXMvZGVmYXVsdHJ1bGVzL3BsLnB5) | `87.50% <ø> (ø)` | |
| [simplemma/token\_sampler.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3Rva2VuX3NhbXBsZXIucHk=) | `93.18% <88.23%> (-3.88%)` | :arrow_down: |
| [...tegies/fallback/lemmatization\_fallback\_strategy.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3N0cmF0ZWdpZXMvZmFsbGJhY2svbGVtbWF0aXphdGlvbl9mYWxsYmFja19zdHJhdGVneS5weQ==) | `88.88% <88.88%> (ø)` | |
| [simplemma/dictionary\_factory.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL2RpY3Rpb25hcnlfZmFjdG9yeS5weQ==) | `95.83% <90.00%> (-1.61%)` | :arrow_down: |
| [simplemma/strategies/lemmatization\_strategy.py](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi#diff-c2ltcGxlbW1hL3N0cmF0ZWdpZXMvbGVtbWF0aXphdGlvbl9zdHJhdGVneS5weQ==) | `90.00% <90.00%> (ø)` | |
| ... and [19 more](https://app.codecov.io/gh/adbar/simplemma/pull/84?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Adrien+Barbaresi) | |


adbar: The merging sequence would be #78, then this one, then #85, correct?

Could you please make sure that all the classes use slots? It can make the code run faster at no cost.
adbar: There are unused imports in dictionary_factory.py and tokenizer.py: ABC.
adbar: Since these are your changes I think it's best if you resolve the merge conflicts.
adbar: Thanks for the changes!

What about putting these lines in the `__init__.py` file to avoid duplicate code?
```
import sys

if sys.version_info >= (3, 8):
    from typing import Protocol
else:
    from typing_extensions import Protocol
```

Patch:
diff --git a/simplemma/__init__.py b/simplemma/__init__.py
index 5b12946..74510e6 100644
--- a/simplemma/__init__.py
+++ b/simplemma/__init__.py
@@ -9,11 +9,12 @@ __version__ = "0.9.1"
 
 from .language_detector import LanguageDetector, in_target_language, langdetect
 from .lemmatizer import Lemmatizer, lemmatize, lemma_iterator, text_lemmatizer, is_known
-from .tokenizer import Tokenizer, simple_tokenizer
+from .tokenizer import Tokenizer, RegexTokenizer, simple_tokenizer
 from .token_sampler import (
     TokenSampler,
+    BaseTokenSampler,
     MostCommonTokenSampler,
     RelaxedMostCommonTokenSampler,
 )
-from .dictionary_factory import DictionaryFactory
+from .dictionary_factory import DictionaryFactory, DefaultDictionaryFactory
 from .dictionary_pickler import *
diff --git a/simplemma/dictionary_factory.py b/simplemma/dictionary_factory.py
index 4b61892..4a8024e 100644
--- a/simplemma/dictionary_factory.py
+++ b/simplemma/dictionary_factory.py
@@ -1,14 +1,21 @@
-"""Parts related to dictonaries."""
+"""DictionaryFactory module. A DictionaryFactory is a class that provides dictionary data."""
+
 import lzma
 import logging
 import pickle
-
+import sys
+from abc import ABC, abstractmethod
 from os import listdir, path
 from pathlib import Path
 from typing import Any, Dict, Optional, Tuple, Union
 
 from functools import lru_cache
 
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
 LOGGER = logging.getLogger(__name__)
 
 
@@ -39,7 +46,16 @@ def _load_dictionary_from_disk(langcode: str) -> Dict[str, str]:
         return pickled_dict
 
 
-class DictionaryFactory:
+class DictionaryFactory(Protocol):
+    @abstractmethod
+    def get_dictionaries(
+        self,
+        langs: Optional[Union[str, Tuple[str, ...]]] = None,
+    ) -> Dict[str, Dict[str, str]]:
+        raise NotImplementedError
+
+
+class DefaultDictionaryFactory(DictionaryFactory):
     __slots__ = ["_data", "_load_dictionary_from_disk"]
 
     def __init__(self, cache_max_size: int = 8):
diff --git a/simplemma/language_detector.py b/simplemma/language_detector.py
index afb1ede..7c4763e 100644
--- a/simplemma/language_detector.py
+++ b/simplemma/language_detector.py
@@ -1,4 +1,4 @@
-"""Experimental language detection."""
+"""LanguageDetector can extract information about languages present in a text like the proportion of text in each language or the main language of the text."""
 
 from operator import itemgetter
 from typing import Dict, List, Tuple, Union
@@ -6,7 +6,7 @@ from typing import Dict, List, Tuple, Union
 from .strategies.lemmatization_strategy import LemmatizationStrategy
 from .strategies.default import DefaultStrategy
 
-from .dictionary_factory import DictionaryFactory
+from .dictionary_factory import DictionaryFactory, DefaultDictionaryFactory
 from .token_sampler import (
     TokenSampler,
     MostCommonTokenSampler,
@@ -18,7 +18,7 @@ def in_target_language(
     text: str,
     lang: Union[str, Tuple[str, ...]],
     greedy: bool = False,
-    dictionary_factory: DictionaryFactory = DictionaryFactory(),
+    dictionary_factory: DictionaryFactory = DefaultDictionaryFactory(),
     token_sampler: TokenSampler = MostCommonTokenSampler(),
 ) -> float:
     """Determine which proportion of the text is in the target language(s)."""
@@ -31,7 +31,7 @@ def langdetect(
     text: str,
     lang: Union[str, Tuple[str, ...]],
     greedy: bool = False,
-    dictionary_factory: DictionaryFactory = DictionaryFactory(),
+    dictionary_factory: DictionaryFactory = DefaultDictionaryFactory(),
     token_samplers: List[TokenSampler] = [
         MostCommonTokenSampler(),
         RelaxedMostCommonTokenSampler(),
@@ -76,7 +76,7 @@ class LanguageDetector:
     def __init__(
         self,
         lang: Union[str, Tuple[str, ...]],
-        dictionary_factory: DictionaryFactory = DictionaryFactory(),
+        dictionary_factory: DictionaryFactory = DefaultDictionaryFactory(),
         token_sampler: TokenSampler = MostCommonTokenSampler(),
         lemmatization_strategy: LemmatizationStrategy = DefaultStrategy(),
     ) -> None:
diff --git a/simplemma/lemmatizer.py b/simplemma/lemmatizer.py
index fbe77c1..b97c8f3 100644
--- a/simplemma/lemmatizer.py
+++ b/simplemma/lemmatizer.py
@@ -1,4 +1,5 @@
-"""Main module."""
+"""Lemmatizer can find the lemma of a token using a search strategy.
+Lemmatizer can also find the tokens in a text and return the list of lemmas."""
 
 from functools import lru_cache
 from typing import Any, List, Iterator, Tuple, Union
@@ -12,8 +13,8 @@ from .strategies.fallback.lemmatization_fallback_strategy import (
 )
 from .strategies.fallback.default import DefaultFallbackStrategy
 
-from .dictionary_factory import DictionaryFactory
-from .tokenizer import Tokenizer
+from .dictionary_factory import DictionaryFactory, DefaultDictionaryFactory
+from .tokenizer import Tokenizer, RegexTokenizer
 
 PUNCTUATION = {".", "?", "!", "…", "¿", "¡"}
 
@@ -29,7 +30,7 @@ def _control_input_type(token: Any) -> None:
 def is_known(
     token: str,
     lang: Union[str, Tuple[str, ...]],
-    dictionary_factory: DictionaryFactory = DictionaryFactory(),
+    dictionary_factory: DictionaryFactory = DefaultDictionaryFactory(),
 ) -> bool:
     return Lemmatizer(dictionary_factory=dictionary_factory).is_known(token, lang)
 
@@ -38,7 +39,7 @@ def lemmatize(
     token: str,
     lang: Union[str, Tuple[str, ...]],
     greedy: bool = False,
-    dictionary_factory: DictionaryFactory = DictionaryFactory(),
+    dictionary_factory: DictionaryFactory = DefaultDictionaryFactory(),
 ) -> str:
     return Lemmatizer(
         dictionary_factory=dictionary_factory,
@@ -50,8 +51,8 @@ def text_lemmatizer(
     text: str,
     lang: Union[str, Tuple[str, ...]],
     greedy: bool = False,
-    dictionary_factory: DictionaryFactory = DictionaryFactory(),
-    tokenizer: Tokenizer = Tokenizer(),
+    dictionary_factory: DictionaryFactory = DefaultDictionaryFactory(),
+    tokenizer: Tokenizer = RegexTokenizer(),
 ) -> List[str]:
     return list(
         lemma_iterator(
@@ -68,8 +69,8 @@ def lemma_iterator(
     text: str,
     lang: Union[str, Tuple[str, ...]],
     greedy: bool = False,
-    dictionary_factory: DictionaryFactory = DictionaryFactory(),
-    tokenizer: Tokenizer = Tokenizer(),
+    dictionary_factory: DictionaryFactory = DefaultDictionaryFactory(),
+    tokenizer: Tokenizer = RegexTokenizer(),
 ) -> Iterator[str]:
     return Lemmatizer(
         dictionary_factory=dictionary_factory,
@@ -90,8 +91,8 @@ class Lemmatizer:
     def __init__(
         self,
         cache_max_size: int = 1048576,
-        dictionary_factory: DictionaryFactory = DictionaryFactory(),
-        tokenizer: Tokenizer = Tokenizer(),
+        dictionary_factory: DictionaryFactory = DefaultDictionaryFactory(),
+        tokenizer: Tokenizer = RegexTokenizer(),
         lemmatization_strategy: LemmatizationStrategy = DefaultStrategy(),
         fallback_lemmatization_strategy: LemmatizationFallbackStrategy = DefaultFallbackStrategy(),
     ) -> None:
diff --git a/simplemma/strategies/fallback/lemmatization_fallback_strategy.py b/simplemma/strategies/fallback/lemmatization_fallback_strategy.py
index a26ea95..b920936 100644
--- a/simplemma/strategies/fallback/lemmatization_fallback_strategy.py
+++ b/simplemma/strategies/fallback/lemmatization_fallback_strategy.py
@@ -1,7 +1,13 @@
-from typing import Dict, Optional
-from abc import ABC
+import sys
+from abc import abstractmethod
 
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
 
-class LemmatizationFallbackStrategy(ABC):
+
+class LemmatizationFallbackStrategy(Protocol):
+    @abstractmethod
     def get_lemma(self, token: str, lang: str) -> str:
         raise NotImplementedError()
diff --git a/simplemma/strategies/lemmatization_strategy.py b/simplemma/strategies/lemmatization_strategy.py
index bf64111..89f1656 100644
--- a/simplemma/strategies/lemmatization_strategy.py
+++ b/simplemma/strategies/lemmatization_strategy.py
@@ -1,8 +1,15 @@
+import sys
 from typing import Dict, Optional
-from abc import ABC
+from abc import abstractmethod
 
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
 
-class LemmatizationStrategy(ABC):
+
+class LemmatizationStrategy(Protocol):
+    @abstractmethod
     def get_lemma(
         self, token: str, lang: str, dictionary: Dict[str, str]
     ) -> Optional[str]:
diff --git a/simplemma/token_sampler.py b/simplemma/token_sampler.py
index c228f12..2cd97a9 100644
--- a/simplemma/token_sampler.py
+++ b/simplemma/token_sampler.py
@@ -1,36 +1,53 @@
-import re
+"""TokenSampler module. A TokenSampler is a class that select samples from a text or a tokens collection."""
 
-from abc import ABC
+import re
+import sys
+from abc import ABC, abstractmethod
 from typing import Iterable, List
 from collections import Counter
-from .tokenizer import Tokenizer
+from .tokenizer import Tokenizer, RegexTokenizer
 
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
 SPLIT_INPUT = re.compile(r"[^\W\d_]{3,}")
 RELAXED_SPLIT_INPUT = re.compile(r"[\w-]{3,}")
 
 
-class TokenSampler(ABC):
+class TokenSampler(Protocol):
+    @abstractmethod
+    def sample_text(self, text: str) -> List[str]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def sample_tokens(self, tokens: Iterable[str]) -> List[str]:
+        raise NotImplementedError
+
+
+class BaseTokenSampler(ABC, TokenSampler):
     __slots__ = ["tokenizer"]
 
     def __init__(
         self,
-        tokenizer: Tokenizer = Tokenizer(SPLIT_INPUT),
+        tokenizer: Tokenizer = RegexTokenizer(SPLIT_INPUT),
     ) -> None:
         self.tokenizer = tokenizer
 
     def sample_text(self, text: str) -> List[str]:
         return self.sample_tokens(self.tokenizer.split_text(text))
 
+    @abstractmethod
     def sample_tokens(self, tokens: Iterable[str]) -> List[str]:
         raise NotImplementedError
 
 
-class MostCommonTokenSampler(TokenSampler):
+class MostCommonTokenSampler(BaseTokenSampler):
     __slots__ = ["capitalized_threshold", "sample_size"]
 
     def __init__(
         self,
-        tokenizer: Tokenizer = Tokenizer(SPLIT_INPUT),
+        tokenizer: Tokenizer = RegexTokenizer(SPLIT_INPUT),
         sample_size: int = 100,
         capitalized_threshold: float = 0.8,
     ) -> None:
@@ -38,9 +55,6 @@ class MostCommonTokenSampler(TokenSampler):
         self.sample_size = sample_size
         self.capitalized_threshold = capitalized_threshold
 
-    def sample_text(self, text: str) -> List[str]:
-        return self.sample_tokens(self.tokenizer.split_text(text))
-
     def sample_tokens(self, tokens: Iterable[str]) -> List[str]:
         """Extract potential tokens, scramble them, potentially get rid of capitalized
         ones, and return the most frequent."""
@@ -59,7 +73,7 @@ class MostCommonTokenSampler(TokenSampler):
 class RelaxedMostCommonTokenSampler(MostCommonTokenSampler):
     def __init__(
         self,
-        tokenizer: Tokenizer = Tokenizer(RELAXED_SPLIT_INPUT),
+        tokenizer: Tokenizer = RegexTokenizer(RELAXED_SPLIT_INPUT),
         sample_size: int = 1000,
         capitalized_threshold: float = 0,
     ) -> None:
diff --git a/simplemma/tokenizer.py b/simplemma/tokenizer.py
index 11cddba..bcb0aa2 100644
--- a/simplemma/tokenizer.py
+++ b/simplemma/tokenizer.py
@@ -1,9 +1,15 @@
-"""Parts related to tokenization."""
+"""Tokenizers module. A Tokenizer is a class that breaks a string into lemmas."""
 
 import re
+import sys
 
+from abc import ABC, abstractmethod
 from typing import Iterator, List, Pattern
 
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
 
 TOKREGEX = re.compile(
     r"(?:"
@@ -22,7 +28,13 @@ def simple_tokenizer(text: str, splitting_regex: Pattern[str] = TOKREGEX) -> Lis
     return splitting_regex.findall(text)
 
 
-class Tokenizer:
+class Tokenizer(Protocol):
+    @abstractmethod
+    def split_text(self, text: str) -> Iterator[str]:
+        raise NotImplementedError
+
+
+class RegexTokenizer(Tokenizer):
     __slots__ = ["splitting_regex"]
 
     def __init__(self, splitting_regex: Pattern[str] = TOKREGEX) -> None:


Test patch:
diff --git a/tests/strategies/test_strategies.py b/tests/strategies/test_strategies.py
index aef8bfd..5bfc4d3 100644
--- a/tests/strategies/test_strategies.py
+++ b/tests/strategies/test_strategies.py
@@ -1,4 +1,4 @@
-from simplemma.dictionary_factory import DictionaryFactory
+from simplemma.dictionary_factory import DefaultDictionaryFactory
 from simplemma.strategies.dictionary_lookup import DictionaryLookupStrategy
 from simplemma.strategies.hyphen_removal import HyphenRemovalStrategy
 from simplemma.strategies.default import DefaultStrategy
@@ -8,7 +8,7 @@ from simplemma.strategies.affix_decomposition import AffixDecompositionStrategy
 
 def test_search() -> None:
     """Test simple and greedy dict search."""
-    dictionary_factory = DictionaryFactory()
+    dictionary_factory = DefaultDictionaryFactory()
     dictionaries = dictionary_factory.get_dictionaries(("en",))
     enDict = dictionaries["en"]
     assert DictionaryLookupStrategy().get_lemma("ignorant", "en", enDict) == "ignorant"
diff --git a/tests/test_dictionary_factory.py b/tests/test_dictionary_factory.py
index 9f7d1dd..9156700 100644
--- a/tests/test_dictionary_factory.py
+++ b/tests/test_dictionary_factory.py
@@ -1,9 +1,9 @@
-from simplemma.dictionary_factory import DictionaryFactory
+from simplemma.dictionary_factory import DefaultDictionaryFactory
 
 
 def test_dictionary_cache() -> None:
     iterations = 10
-    dictionaries = DictionaryFactory()
+    dictionaries = DefaultDictionaryFactory()
     for _ in range(iterations):
         dictionaries.get_dictionaries("en")
         dictionaries.get_dictionaries("de")
diff --git a/tests/test_lemmatizer.py b/tests/test_lemmatizer.py
index 444a04f..b7043cc 100644
--- a/tests/test_lemmatizer.py
+++ b/tests/test_lemmatizer.py
@@ -5,7 +5,7 @@ import pytest
 from typing import Dict, Optional, Tuple, Union
 
 import simplemma
-from simplemma import lemmatize, Lemmatizer, DictionaryFactory
+from simplemma import lemmatize, Lemmatizer, DictionaryFactory, DefaultDictionaryFactory
 from simplemma.strategies.fallback.raise_error import RaiseErrorFallbackStrategy
 
 logging.basicConfig(level=logging.DEBUG)
@@ -81,7 +81,7 @@ def test_readme() -> None:
 def test_logic() -> None:
     """Test if certain code parts correspond to the intended logic."""
     # missing languages or faulty language codes
-    dictionary_factory = DictionaryFactory()
+    dictionary_factory = DefaultDictionaryFactory()
     with pytest.raises(ValueError):
         dictionaries = dictionary_factory.get_dictionaries(("abc"))
 
diff --git a/tests/test_token_sampler.py b/tests/test_token_sampler.py
index cdb2c0f..0de9d8f 100644
--- a/tests/test_token_sampler.py
+++ b/tests/test_token_sampler.py
@@ -1,13 +1,13 @@
 from typing import Iterable, List
 
 from simplemma.token_sampler import (
-    TokenSampler,
+    BaseTokenSampler,
     MostCommonTokenSampler,
     RelaxedMostCommonTokenSampler,
 )
 
 
-class CustomTokenSampler(TokenSampler):
+class CustomTokenSampler(BaseTokenSampler):
     def __init__(self, skip_tokens: int) -> None:
         super().__init__()
         self.skip_tokens: int = skip_tokens