Add new OP: token_num_filter

app.py

@@ -211,6 +211,7 @@ def get_min_max_step(data):
     'special_characters_filter': [StatsKeys.special_char_ratio],
     'stopwords_filter': [StatsKeys.stopwords_ratio],
     'text_length_filter': [StatsKeys.text_len],
+    'token_num_filter': [StatsKeys.num_token],
     'words_num_filter': [StatsKeys.num_words],
     'word_repetition_filter': [StatsKeys.word_rep_ratio],
 }

configs/config_all.yaml

@@ -13,8 +13,8 @@ np: 4                                                       # number of subproce
 text_keys: 'content'                                        # the key name of field where the sample texts to be processed, e.g., `text`, `instruction`, `output`, ...
                                                             # Note: currently, we support specify only ONE key for each op, for cases requiring multiple keys, users can specify the op multiple times. We will only use the first key of `text_keys` when you set multiple keys.
 suffixes: []                                                # the suffix of files that will be read. For example: '.txt', 'txt' or ['txt', '.pdf', 'docx']
-use_cache: true                                             # whether to use the cache management of hugging face datasets. It might take up lots of disk space when using cache
-ds_cache_dir: '~/.cache/huggingface/datasets'               # cache dir for hugging face datasets. In default it's the default cache dir "~/.cache/huggingface/datasets". If this argument is reset by users, it will override the default cache dir
+use_cache: true                                             # whether to use the cache management of Hugging Face datasets. It might take up lots of disk space when using cache
+ds_cache_dir: '~/.cache/huggingface/datasets'               # cache dir for Hugging Face datasets. In default it's the default cache dir "~/.cache/huggingface/datasets". If this argument is reset by users, it will override the default cache dir
 use_checkpoint: false                                       # whether to use the checkpoint management to save the latest version of dataset to work dir when processing. Rerun the same config will reload the checkpoint and skip ops before it. Cache will be disabled when using checkpoint. If args of ops before the checkpoint are changed, all ops will be rerun from the beginning.
 temp_dir: null                                              # the path to the temp directory to store intermediate caches when cache is disabled, these cache files will be removed on-the-fly. In default, it's None, so the temp dir will be specified by system. NOTICE: you should be caution when setting this argument because it might cause unexpected program behaviors when this path is set to an unsafe directory.
 open_tracer: false                                          # whether to open the tracer to trace the changes during process. It might take more time when opening tracer
@@ -127,6 +127,10 @@ process:
   - text_length_filter:                                     # filter text with length out of specific range
       min_len: 10                                             # the min length of filter range
       max_len: 10000                                          # the max length of filter range
+  - token_num_filter:                                       # filter text with total token number out of specific range
+      hf_tokenizer: EleutherAI/pythia-6.9b-deduped            # name of used Hugging Face tokenizer
+      min_num: 10                                             # the min number of filter range
+      max_num: 10000                                          # the max number of filter range
   - words_num_filter:                                       # filter text with number of words out of specific range
       lang: en                                                # sample in which language
       tokenization: false                                     # whether to use model to tokenize documents

data_juicer/ops/filter/__init__.py

@@ -4,4 +4,4 @@
                perplexity_filter, special_characters_filter,
                specified_field_filter, specified_numeric_field_filter,
                stopwords_filter, suffix_filter, text_length_filter,
-               word_num_filter, word_repetition_filter)
+               token_num_filter, word_num_filter, word_repetition_filter)

data_juicer/ops/filter/token_num_filter.py

@@ -0,0 +1,61 @@
+import sys
+
+from jsonargparse.typing import PositiveInt
+
+from data_juicer.utils.constant import Fields, StatsKeys
+from data_juicer.utils.model_utils import prepare_model, get_model
+
+from ..base_op import OPERATORS, Filter
+from ..common import get_words_from_document
+
+
+@OPERATORS.register_module('token_num_filter')
+class TokenNumFilter(Filter):
+    """Filter to keep samples with total token number within a specific
+    range."""
+
+    def __init__(self,
+                 hf_tokenizer: str = 'EleutherAI/pythia-6.9b-deduped',
+                 min_num: PositiveInt = 10,
+                 max_num: PositiveInt = sys.maxsize,
+                 *args,
+                 **kwargs):
+        """
+        Initialization method.
+
+        :param hf_tokenizer: the tokenizer name of Hugging Face tokenizers.
+        :param min_num: The min filter token number in this op, samples
+            will be filtered if their token number is below this
+            parameter.
+        :param max_num: The max filter token number in this op, samples
+            will be filtered if their token number exceeds this
+            parameter.
+        :param args: extra args
+        :param kwargs: extra args
+        """
+        super().__init__(*args, **kwargs)
+        self.min_num = min_num
+        self.max_num = max_num
+        self.hf_tokenizer = hf_tokenizer
+        self.model_key = prepare_model(model_type='huggingface',
+                                       model_key=hf_tokenizer)
+
+    def compute_stats(self, sample):
+        # check if it's computed already
+        if StatsKeys.num_token in sample[Fields.stats]:
+            return sample
+
+        tokenizer = get_model(self.model_key, model_type='huggingface')
+        tokens = get_words_from_document(
+            sample[self.text_key],
+            token_func=tokenizer.tokenize if tokenizer else None
+        )
+        sample[Fields.stats][StatsKeys.num_token] = len(tokens)
+        return sample
+
+    def process(self, sample):
+        if self.min_num <= sample[Fields.stats][
+                StatsKeys.num_token] <= self.max_num:
+            return True
+        else:
+            return False

data_juicer/utils/constant.py

@@ -21,6 +21,7 @@ class StatsKeys(object):
     special_char_ratio = 'special_char_ratio'
     stopwords_ratio = 'stopwords_ratio'
     text_len = 'text_len'
+    num_token = 'num_token'
     num_words = 'num_words'
     word_rep_ratio = 'word_rep_ratio'
 

data_juicer/utils/model_utils.py

@@ -61,12 +61,12 @@ def check_model(model_name, args=(), force=False):
 
         try:
             model_link = os.path.join(MODEL_LINKS, true_model_name)
-            wget.download(model_link, mdp)
+            wget.download(model_link, mdp, bar=None)
         except:  # noqa: E722
             try:
                 backup_model_link = os.path.join(
                     BACKUP_MODEL_LINKS[model_name], true_model_name)
-                wget.download(backup_model_link, mdp)
+                wget.download(backup_model_link, mdp, bar=None)
             except:  # noqa: E722
                 logger.error(
                     f'Downloading model [{true_model_name}] error. '
@@ -165,7 +165,8 @@ def prepare_huggingface_tokenizer(tokenizer_name):
     """
     from transformers import AutoTokenizer
     logger.info('Loading tokenizer from HuggingFace...')
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name,
+                                              trust_remote_code=True)
     return tokenizer
 
 def prepare_diversity_model(model_name, lang):
@@ -178,7 +179,6 @@ def prepare_diversity_model(model_name, lang):
     :return: corresponding diversity model
     """
     import spacy
-    print(lang)
     assert lang in ['zh', 'en'], 'Diversity only support zh and en'
     model_name = model_name % lang
     logger.info(f'Loading spacy model [{model_name}]...')

demos/data_visualization_op_effect/app.py

@@ -147,6 +147,7 @@ def get_min_max_step(data):
     'special_characters_filter': [StatsKeys.special_char_ratio],
     'stopwords_filter': [StatsKeys.stopwords_ratio],
     'text_length_filter': [StatsKeys.text_len],
+    'token_num_filter': [StatsKeys.num_token],
     'words_num_filter': [StatsKeys.num_words],
     'word_repetition_filter': [StatsKeys.word_rep_ratio],
 }

demos/tool_quality_classifier/quality_classifier/qc_utils.py

@@ -54,7 +54,8 @@ def prepare_model(model_name, model_path=DATA_JUICER_MODELS_CACHE):
     # No specific models in local file systems. Download them from remote.
     os.makedirs(model_path, exist_ok=True)
     wget.download(os.path.join(MODEL_LINKS, f'{model_name}.zip'),
-                  os.path.join(model_path, f'{model_name}.zip'))
+                  os.path.join(model_path, f'{model_name}.zip'),
+                  bar=None)
     with zipfile.ZipFile(os.path.join(model_path, f'{model_name}.zip')) as zip:
         zip.extractall(os.path.join(model_path))
     return PipelineModel.load(real_model_path)

docs/Operators.md

@@ -11,7 +11,7 @@ The operators in Data-Juicer are categorized into 5 types.
 |-----------------------------------|:------:|-------------------------------------------------|
 | [ Formatter ]( #formatter )       |   7    | Discovers, loads, and canonicalizes source data |
 | [ Mapper ]( #mapper )             |   19   | Edits and transforms samples                    |
-| [ Filter ]( #filter )             |   15   | Filters out low-quality samples                 |
+| [ Filter ]( #filter )             |   16   | Filters out low-quality samples                 |
 | [ Deduplicator ]( #deduplicator ) |   3    | Detects and removes duplicate samples           |
 | [ Selector ]( #selector )         |   2    | Selects top samples based on ranking            |
 
@@ -83,6 +83,7 @@ All the specific operators are listed below, each featured with several capabili
 | stopwords_filter               | General | en, zh | Keeps samples with stopword ratio above the specified threshold                            |
 | suffix_filter                  | General | en, zh | Keeps samples with specified suffixes                                                      |
 | text_length_filter             | General | en, zh | Keeps samples with total text length within the specified range                            |
+| token_num_filter               | General | en, zh | Keeps samples with token count within the specified range                                  |
 | word_num_filter                | General | en, zh | Keeps samples with word count within the specified range                                   |
 | word_repetition_filter         | General | en, zh | Keeps samples with word-level n-gram repetition ratio within the specified range           |
 

docs/Operators_ZH.md

@@ -6,13 +6,13 @@
 
 Data-Juicer 中的算子分为以下 5 种类型。
 
-| 类型                                 | 数量  | 描述            |
-|------------------------------------|:---:|---------------|
-| [ Formatter ]( #formatter )        |  7  | 发现、加载、规范化原始数据 |
-| [ Mapper ]( #mapper )              | 19  | 对数据样本进行编辑和转换  |
-| [ Filter ]( #filter )              | 15  | 过滤低质量样本       |
-| [ Deduplicator ]( #deduplicator )  |  3  | 识别、删除重复样本     |
-| [ Selector ]( #selector )          |  2  | 基于排序选取高质量样本   |
+| 类型                                 | 数量 | 描述            |
+|------------------------------------|:--:|---------------|
+| [ Formatter ]( #formatter )        |  7 | 发现、加载、规范化原始数据 |
+| [ Mapper ]( #mapper )              | 19 | 对数据样本进行编辑和转换  |
+| [ Filter ]( #filter )              | 16  | 过滤低质量样本       |
+| [ Deduplicator ]( #deduplicator )  |  3 | 识别、删除重复样本     |
+| [ Selector ]( #selector )          |  2 | 基于排序选取高质量样本   |
 
 下面列出所有具体算子，每种算子都通过多个标签来注明其主要功能。
 
@@ -64,23 +64,24 @@ Data-Juicer 中的算子分为以下 5 种类型。
 
 ## Filter <a name="filter"/>
 
-| 算子                              | 场景       | 语言      | 描述                                                       |
-|---------------------------------|----------|---------|----------------------------------------------------------|
-| alphanumeric_filter             | General  | en, zh  | 保留字母数字比例在指定范围内的样本                                        |
-| average_line_length_filter      | Code     | en, zh  | 保留平均行长度在指定范围内的样本                                         |
-| character_repetition_filter     | General  | en, zh  | 保留 char-level n-gram 重复比率在指定范围内的样本                       |
-| flagged_words_filter            | General  | en, zh  | 保留使标记字比率保持在指定阈值以下的样本                                     |
-| language_id_score_filter        | General  | en, zh  | 保留特定语言的样本，通过预测的置信度得分来判断                                  |
-| maximum_line_length_filter      | Code     | en, zh  | 保留最大行长度在指定范围内的样本                                         |
-| perplexity_filter               | General  | en, zh  | 保留困惑度低于指定阈值的样本                                           |
-| special_characters_filter       | General  | en, zh  | 保留 special-char 比率的在指定范围内的样本                             |
-| specified_field_filter          | General  | en, zh  | 根据字段过滤样本，要求字段的值处于指定目标中                                   |
-| specified_numeric_field_filter  | General  | en, zh  | 根据字段过滤样本，要求字段的值处于指定范围（针对数字类型）                            |
-| stopwords_filter                | General  | en, zh  | 保留停用词比率高于指定阈值的样本                                         |
-| suffix_filter                   | General  | en, zh  | 保留包含特定后缀的样本                                              |
-| text_length_filter              | General  | en, zh  | 保留总文本长度在指定范围内的样本                                         |
-| word_num_filter                 | General  | en, zh  | 保留字数在指定范围内的样本                                            |
-| word_repetition_filter          | General  | en, zh  | 保留 word-level n-gram 重复比率在指定范围内的样本                       |
+| 算子                             | 场景       | 语言      | 描述                                 |
+|--------------------------------|----------|---------|------------------------------------|
+| alphanumeric_filter            | General  | en, zh  | 保留字母数字比例在指定范围内的样本                  |
+| average_line_length_filter     | Code     | en, zh  | 保留平均行长度在指定范围内的样本                   |
+| character_repetition_filter    | General  | en, zh  | 保留 char-level n-gram 重复比率在指定范围内的样本 |
+| flagged_words_filter           | General  | en, zh  | 保留使标记字比率保持在指定阈值以下的样本               |
+| language_id_score_filter       | General  | en, zh  | 保留特定语言的样本，通过预测的置信度得分来判断            |
+| maximum_line_length_filter     | Code     | en, zh  | 保留最大行长度在指定范围内的样本                   |
+| perplexity_filter              | General  | en, zh  | 保留困惑度低于指定阈值的样本                     |
+| special_characters_filter      | General  | en, zh  | 保留 special-char 比率的在指定范围内的样本       |
+| specified_field_filter         | General  | en, zh  | 根据字段过滤样本，要求字段的值处于指定目标中             |
+| specified_numeric_field_filter | General  | en, zh  | 根据字段过滤样本，要求字段的值处于指定范围（针对数字类型）      |
+| stopwords_filter               | General  | en, zh  | 保留停用词比率高于指定阈值的样本                   |
+| suffix_filter                  | General  | en, zh  | 保留包含特定后缀的样本                        |
+| text_length_filter             | General  | en, zh  | 保留总文本长度在指定范围内的样本                   |
+| token_num_filter               | General  | en, zh  | 保留token数在指定范围内的样本                  |
+| word_num_filter                | General  | en, zh  | 保留字数在指定范围内的样本                      |
+| word_repetition_filter         | General  | en, zh  | 保留 word-level n-gram 重复比率在指定范围内的样本 |
 
 ## Deduplicator <a name="deduplicator"/>
 

tests/ops/filter/test_token_num_filter.py

@@ -0,0 +1,56 @@
+import unittest
+
+from datasets import Dataset
+
+from data_juicer.ops.filter.token_num_filter import TokenNumFilter
+from data_juicer.utils.constant import Fields, StatsKeys
+
+
+class WordNumFilterTest(unittest.TestCase):
+
+    def test_token_num(self):
+        src = [
+            {"text": "Today is Sunday and it's a happy day!"},
+            {"text": "Do you need a cup of coffee?"},
+            {"text": "你好，请问你是谁"},
+            {"text": "Sur la plateforme MT4, plusieurs manières d'accéder à "
+                     "ces fonctionnalités sont conçues simultanément."},
+            {"text": "欢迎来到阿里巴巴！"},
+            {"text": "This paper proposed a novel method on LLM pretraining."},
+        ]
+        tgt = [
+            10, 8, 9, 31, 14, 12,
+        ]
+        ds = Dataset.from_list(src)
+        op = TokenNumFilter()
+        ds = ds.add_column(Fields.stats, [{}] * len(ds))
+        ds = ds.map(op.compute_stats)
+        stats = ds[Fields.stats]
+        self.assertEqual([item[StatsKeys.num_token] for item in stats], tgt)
+
+    def test_token_num_filter(self):
+        src = [
+            {"text": "Today is Sunday and it's a happy day!"},
+            {"text": "Do you need a cup of coffee?"},
+            {"text": "你好，请问你是谁"},
+            {"text": "Sur la plateforme MT4, plusieurs manières d'accéder à "
+                     "ces fonctionnalités sont conçues simultanément."},
+            {"text": "欢迎来到阿里巴巴！"},
+            {"text": "This paper proposed a novel method on LLM pretraining."},
+        ]
+        tgt = [
+            {"text": "Today is Sunday and it's a happy day!"},
+            {"text": "欢迎来到阿里巴巴！"},
+            {"text": "This paper proposed a novel method on LLM pretraining."},
+        ]
+        ds = Dataset.from_list(src)
+        tgt = Dataset.from_list(tgt)
+        op = TokenNumFilter(min_num=10, max_num=20)
+        ds = ds.add_column(Fields.stats, [{}] * len(ds))
+        ds = ds.map(op.compute_stats)
+        ds = ds.filter(op.process)
+        self.assertEqual(ds['text'], tgt['text'])
+
+
+if __name__ == '__main__':
+    unittest.main()

tools/postprocess/README.md

@@ -4,6 +4,26 @@ This folder contains some postprocess scripts for additional processing of your
 
 ## Usage
 
+### Count tokens for datasets
+
+Use [count_token.py](count_token.py) to count tokens for datasets.
+
+```shell
+python tools/postprocess/count_token.py        \
+    --data_path            <data_path>         \
+    --text_keys            <text_keys>         \
+    --tokenizer_method     <tokenizer_method>  \
+    --num_proc             <num_proc>
+
+# get help
+python tools/postprocess/count_token.py --help
+```
+
+- `data_path`: path to the input dataset. Only support `jsonl` now.
+- `text_keys`: field keys that will be considered into token counts.
+- `tokenizer_method`: name of the Hugging Face tokenizer.
+- `num_proc`: number of processes to count tokens.
+
 ### Mix multiple datasets with optional weights
 
 Use [data_mixture.py](data_mixture.py) to mix multiple datasets.