[text] huggingface tokenizer

test/wenet/text/test_hugging_face_tokenizer.py

@@ -0,0 +1,96 @@
+import os
+import pytest
+
+from wenet.text.hugging_face_tokenizer import HuggingFaceTokenizer
+
+try:
+    import transformers  # noqa
+except ImportError:
+    os.system('pip install --no-input transformers')
+    import transformers  # noqa
+
+
+@pytest.fixture(params=["bert-base-cased"])
+def hugging_face_tokenizer(request):
+    return HuggingFaceTokenizer(request.param)
+
+
+def test_text2tokens(hugging_face_tokenizer: HuggingFaceTokenizer):
+    tokenizer = hugging_face_tokenizer
+    text = "hello wenet very cool!"
+    expected = ['hello', 'we', '##net', 'very', 'cool', '!']
+    assert all(h == r for h, r in zip(tokenizer.text2tokens(text), expected))
+
+
+def test_tokens2text(hugging_face_tokenizer: HuggingFaceTokenizer):
+    tokenizer = hugging_face_tokenizer
+    inputs = ['hello', 'we', '##net', 'very', 'cool', '!']
+    expected = "hello wenet very cool!"
+
+    result = tokenizer.tokens2text(inputs)
+    assert result == expected
+
+
+def test_tokens2ids(hugging_face_tokenizer: HuggingFaceTokenizer):
+    tokenizer = hugging_face_tokenizer
+    inputs = ['hello', 'we', '##net', 'very', 'cool', '!']
+    expected = [19082, 1195, 6097, 1304, 4348, 106]
+    tokens = tokenizer.tokens2ids(inputs)
+    assert len(tokens) == len(expected)
+    assert all(h == r for (h, r) in zip(tokens, expected))
+
+
+def test_ids2tokens(hugging_face_tokenizer: HuggingFaceTokenizer):
+    tokenizer = hugging_face_tokenizer
+    ids = [19082, 1195, 6097, 1304, 4348, 106]
+    expected = ['hello', 'we', '##net', 'very', 'cool', '!']
+    results = tokenizer.ids2tokens(ids)
+    assert len(results) == len(expected)
+    assert all(h == r for (h, r) in zip(results, expected))
+
+
+def test_tokenize(hugging_face_tokenizer: HuggingFaceTokenizer):
+    tokenizer = hugging_face_tokenizer
+
+    text = "hello wenet very cool!"
+    ids = [19082, 1195, 6097, 1304, 4348, 106]
+    tokens = ['hello', 'we', '##net', 'very', 'cool', '!']
+
+    r_tokens, r_ids = tokenizer.tokenize(text)
+    assert len(r_tokens) == len(tokens)
+    assert all(h == r for (h, r) in zip(r_tokens, tokens))
+    assert len(r_ids) == len(ids)
+    assert all(h == r for (h, r) in zip(r_ids, ids))
+
+
+def test_detokenize(hugging_face_tokenizer: HuggingFaceTokenizer):
+    tokenizer = hugging_face_tokenizer
+    text = "hello wenet very cool!"
+    ids = [19082, 1195, 6097, 1304, 4348, 106]
+    tokens = ['hello', 'we', '##net', 'very', 'cool', '!']
+
+    r_text, r_tokens = tokenizer.detokenize(ids)
+    assert r_text == text
+    assert len(r_tokens) == len(tokens)
+    assert all(h == r for (h, r) in zip(r_tokens, tokens))
+
+
+def test_vocab_size(hugging_face_tokenizer: HuggingFaceTokenizer):
+    assert hugging_face_tokenizer.vocab_size() == 28996
+    assert hugging_face_tokenizer.vocab_size() == len(
+        hugging_face_tokenizer.symbol_table)
+
+
+def test_tongyi_tokenizer():
+    # NOTE(Mddct): tongyi need extra matplotlib package
+    os.system('pip install --no-input matplotlib')
+    model_dir = 'Qwen/Qwen-Audio-Chat'
+    tongyi_tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_dir, trust_remote_code=True)
+    tokenizer = HuggingFaceTokenizer(model_dir, trust_remote_code=True)
+    text = "from transformers import AutoModelForCausalLM, AutoTokenizer"
+    tongyi_result = tongyi_tokenizer.tokenize(text)
+    result, _ = tokenizer.tokenize(text)
+
+    assert len(result) == len(tongyi_result)
+    assert all(h == r for (h, r) in zip(result, tongyi_result))

test/wenet/text/test_parallel.py

@@ -3,6 +3,7 @@
 from wenet.text.base_tokenizer import BaseTokenizer
 
 from wenet.text.bpe_tokenizer import BpeTokenizer
+from wenet.text.hugging_face_tokenizer import HuggingFaceTokenizer
 from wenet.text.whisper_tokenizer import WhisperTokenizer
 
 
@@ -47,7 +48,7 @@ def test_bpe_tokenzier_parallel():
     symbol_table_path = "test/resources/librispeech.words.txt"
     bpe_model = "test/resources/librispeech.train_960_unigram5000.bpemodel"
 
-    inputs = ["WENR IS SIMPLE", "GOOD"]
+    inputs = ["WENT IS SIMPLE", "GOOD"]
     tokenizer = BpeTokenizer(bpe_model, symbol_table_path)
     partial_tokenize = partial(consistency, tokenizer)
     with Pool(processes=len(inputs)) as pool:
@@ -63,7 +64,7 @@ def test_bpe_tokenizer_parallel_after_property():
     symbol_table_path = "test/resources/librispeech.words.txt"
     bpe_model = "test/resources/librispeech.train_960_unigram5000.bpemodel"
 
-    inputs = ["WENR IS SIMPLE", "GOOD"]
+    inputs = ["WENT IS SIMPLE", "GOOD"]
     tokenizer = BpeTokenizer(bpe_model, symbol_table_path)
     _ = tokenizer.vocab_size
     _ = tokenizer.symbol_table
@@ -76,3 +77,20 @@ def test_bpe_tokenizer_parallel_after_property():
     results.sort()
 
     assert all(h == r for (h, r) in zip(results, inputs))
+
+
+def test_hugging_face_tokenizer():
+    tokenizer = HuggingFaceTokenizer("bert-base-cased")
+
+    _ = tokenizer.vocab_size
+    _ = tokenizer.symbol_table
+
+    inputs = ["wenet is simple", "good"]
+    partial_tokenize = partial(consistency, tokenizer)
+    with Pool(processes=len(inputs)) as pool:
+        results = pool.map(partial_tokenize, inputs)
+
+    inputs.sort()
+    results.sort()
+
+    assert all(h == r for (h, r) in zip(results, inputs))

wenet/text/base_tokenizer.py

@@ -1,39 +1,41 @@
 from abc import ABC, abstractmethod, abstractproperty
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union
+
+T = Union[str, bytes]
 
 
 class BaseTokenizer(ABC):
 
-    def tokenize(self, line: str) -> Tuple[List[str], List[int]]:
+    def tokenize(self, line: str) -> Tuple[List[T], List[int]]:
         tokens = self.text2tokens(line)
         ids = self.tokens2ids(tokens)
         return tokens, ids
 
-    def detokenize(self, ids: List[int]) -> Tuple[str, List[str]]:
+    def detokenize(self, ids: List[int]) -> Tuple[str, List[T]]:
         tokens = self.ids2tokens(ids)
         text = self.tokens2text(tokens)
         return text, tokens
 
     @abstractmethod
-    def text2tokens(self, line: str) -> List[str]:
+    def text2tokens(self, line: str) -> List[T]:
         raise NotImplementedError("abstract method")
 
     @abstractmethod
-    def tokens2text(self, tokens: List[str]) -> str:
+    def tokens2text(self, tokens: List[T]) -> str:
         raise NotImplementedError("abstract method")
 
     @abstractmethod
-    def tokens2ids(self, tokens: List[str]) -> List[int]:
+    def tokens2ids(self, tokens: List[T]) -> List[int]:
         raise NotImplementedError("abstract method")
 
     @abstractmethod
-    def ids2tokens(self, ids: List[int]) -> List[str]:
+    def ids2tokens(self, ids: List[int]) -> List[T]:
         raise NotImplementedError("abstract method")
 
     @abstractmethod
     def vocab_size(self) -> int:
         raise NotImplementedError("abstract method")
 
     @abstractproperty
-    def symbol_table(self) -> Dict[str, int]:
+    def symbol_table(self) -> Dict[T, int]:
         raise NotImplementedError("abstract method")

wenet/text/hugging_face_tokenizer.py

@@ -0,0 +1,58 @@
+from os import PathLike
+from typing import Dict, List, Union
+from wenet.text.base_tokenizer import BaseTokenizer, T as Type
+
+
+class HuggingFaceTokenizer(BaseTokenizer):
+
+    def __init__(self, model: Union[str, PathLike], *args, **kwargs) -> None:
+        # NOTE(Mddct): don't build here, pickle issues
+        self.model = model
+        self.tokenizer = None
+
+        self.args = args
+        self.kwargs = kwargs
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state['tokenizer']
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        recovery = {'tokenizer': None}
+        self.__dict__.update(recovery)
+
+    def _build_hugging_face(self):
+        from transformers import AutoTokenizer
+        if self.tokenizer is None:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model, **self.kwargs)
+            self.t2i = self.tokenizer.get_vocab()
+
+    def text2tokens(self, line: str) -> List[Type]:
+        self._build_hugging_face()
+        return self.tokenizer.tokenize(line)
+
+    def tokens2text(self, tokens: List[Type]) -> str:
+        self._build_hugging_face()
+        ids = self.tokens2ids(tokens)
+        return self.tokenizer.decode(ids)
+
+    def tokens2ids(self, tokens: List[Type]) -> List[int]:
+        self._build_hugging_face()
+        return self.tokenizer.convert_tokens_to_ids(tokens)
+
+    def ids2tokens(self, ids: List[int]) -> List[Type]:
+        self._build_hugging_face()
+        return self.tokenizer.convert_ids_to_tokens(ids)
+
+    def vocab_size(self) -> int:
+        self._build_hugging_face()
+        # TODO: we need special tokenize size in future
+        return len(self.tokenizer)
+
+    @property
+    def symbol_table(self) -> Dict[Type, int]:
+        self._build_hugging_face()
+        return self.t2i