We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG]: fine train llama-2-7b-hf prepare data set error , `bos_token` and `eos_token` should be the same with `conversation_template.seps`.
### Is there an existing issue for this bug?

- [X] I have searched the existing issues

### 🐛 Describe the bug

 python prepare_sft_dataset.py --data_input_dirs /models/train-data-dir  --tokenizer_dir /models/llama-2-13b-hf --data_output_dirs /models/train-data-dir/out  --num_spliced_dataset_bins 1




 File "/ColossalAI/applications/Colossal-LLaMA/colossal_llama/dataset/spliced_and_tokenized_dataset.py", line 90, in supervised_tokenize_sft
    tokenizer.bos_token == conversation_template.seps[0] and tokenizer.eos_token == conversation_template.seps[1]
AssertionError: `bos_token` and `eos_token` should be the same with `conversation_template.seps`.
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/ColossalAI/applications/Colossal-LLaMA/prepare_sft_dataset.py", line 146, in <module>
    main()
  File "/ColossalAI/applications/Colossal-LLaMA/prepare_sft_dataset.py", line 102, in main
    dataset = dataset.map(
  File "/usr/local/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 602, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 567, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3248, in map
    for rank, done, content in iflatmap_unordered(
  File "/usr/local/lib/python3.10/site-packages/datasets/utils/py_utils.py", line 718, in iflatmap_unordered
    [async_result.get(timeout=0.05) for async_result in async_results]
  File "/usr/local/lib/python3.10/site-packages/datasets/utils/py_utils.py", line 718, in <listcomp>
    [async_result.get(timeout=0.05) for async_result in async_results]
  File "/usr/local/lib/python3.10/site-packages/multiprocess/pool.py", line 774, in get
    raise self._value
**AssertionError: `bos_token` and `eos_token` should be the same with `conversation_template.seps`.**

### Environment

------------ Environment ------------
Colossal-AI version: 0.3.7
PyTorch version: 2.1.0
System CUDA version: 11.0
CUDA version required by PyTorch: 11.8

tokenizer_config.json 

  "add_bos_token": true,
  "add_eos_token": false,
  "bos_token": {
    "__type": "AddedToken",
    "content": "<s>",
    "lstrip": false,
    "normalized": true,
    "rstrip": false,
    "single_word": false
  },
  "clean_up_tokenization_spaces": false,
  "eos_token": {
    "__type": "AddedToken",
    "content": "</s>",
    "lstrip": false,
    "normalized": true,
    "rstrip": false,
    "single_word": false
  },
[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `applications/Colossal-LLaMA/prepare_sft_dataset.py`
Content:
```
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 """
4 Prepare sft dataset for fine-tuning
5 """
6 
7 import argparse
8 import json
9 import math
10 import os
11 from multiprocessing import cpu_count
12 
13 from colossal_llama.dataset.conversation import default_conversation
14 from colossal_llama.dataset.spliced_and_tokenized_dataset import supervised_tokenize_sft
15 from datasets import dataset_dict, load_dataset
16 from transformers import AddedToken, AutoTokenizer
17 
18 from colossalai.logging import get_dist_logger
19 
20 logger = get_dist_logger()
21 
22 
23 def main():
24     parser = argparse.ArgumentParser()
25     parser.add_argument(
26         "--data_input_dirs",
27         type=str,
28         required=True,
29         default=None,
30         help="Comma(i.e., ',') separated list of all data directories containing `.jsonl` data files.",
31     )
32     parser.add_argument(
33         "--tokenizer_dir", type=str, required=True, default=None, help="A directory containing the tokenizer"
34     )
35     parser.add_argument("--data_output_dirs", type=str, default="data_output_dirs", help="Data output directory")
36     parser.add_argument("--max_length", type=int, default=8192, help="Max length of each spliced tokenized sequence")
37     parser.add_argument("--num_spliced_dataset_bins", type=int, default=10, help="Number of spliced dataset bins")
38     parser.add_argument("--llama_version", type=int, default=3, help="LLaMA version")
39     args = parser.parse_args()
40 
41     if args.num_spliced_dataset_bins >= 100000:
42         raise ValueError("Too many spliced divisions, must be smaller than 100000")
43 
44     args.data_cache_dir = os.path.join(args.data_output_dirs, "cache")
45     args.data_jsonl_output_dir = os.path.join(args.data_output_dirs, "jsonl")
46     args.data_arrow_output_dir = os.path.join(args.data_output_dirs, "arrow")
47 
48     if not os.path.exists(args.data_cache_dir):
49         os.makedirs(args.data_cache_dir)
50     if not os.path.exists(args.data_jsonl_output_dir):
51         os.makedirs(args.data_jsonl_output_dir)
52     if not os.path.exists(args.data_arrow_output_dir):
53         os.makedirs(args.data_arrow_output_dir)
54 
55     # Prepare to all input datasets
56     input_data_paths = []
57     input_data_dirs = args.data_input_dirs.split(",")
58     for ds_dir in input_data_dirs:
59         ds_dir = os.path.abspath(ds_dir)
60         assert os.path.exists(ds_dir), f"Not find data dir {ds_dir}"
61         ds_files = [name for name in os.listdir(ds_dir) if name.endswith(".jsonl")]
62         ds_paths = [os.path.join(ds_dir, name) for name in ds_files]
63         input_data_paths.extend(ds_paths)
64 
65     # Prepare to data splitting.
66     train_splits = []
67     split_interval = math.ceil(100 / args.num_spliced_dataset_bins)
68     for i in range(0, 100, split_interval):
69         start = i
70         end = i + split_interval
71         if end > 100:
72             end = 100
73         train_splits.append(f"train[{start}%:{end}%]")
74 
75     # Prepare to the tokenizer.
76     tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir)
77 
78     # Fix </s> split issue: https://github.com/huggingface/transformers/issues/23833
79     if args.llama_version == 2:
80         tokenizer.add_tokens(AddedToken("</s>", normalized=False, special=True), special_tokens=True)
81 
82     tokenizer.add_bos_token = False
83     tokenizer.add_eos_token = False
84     if tokenizer.pad_token is None:
85         if tokenizer.unk_token is not None:
86             tokenizer.pad_token = tokenizer.unk_token
87         else:
88             tokenizer.pad_token = tokenizer.eos_token
89             tokenizer.unk_token = tokenizer.eos_token
90 
91     list_dataset = load_dataset(
92         path="json",
93         data_files=input_data_paths,
94         cache_dir=os.path.join(args.data_cache_dir, "raw"),
95         keep_in_memory=False,
96         split=train_splits,
97         num_proc=cpu_count(),
98     )
99     for index, dataset in enumerate(list_dataset):
100         assert isinstance(dataset, dataset_dict.Dataset)
101         logger.info(f"Start to process part-{index}/{len(list_dataset)} of all original datasets.")
102         dataset = dataset.map(
103             function=supervised_tokenize_sft,
104             fn_kwargs={
105                 "tokenizer": tokenizer,
106                 "conversation_template": default_conversation,
107                 "max_length": args.max_length,
108             },
109             keep_in_memory=False,
110             num_proc=min(len(dataset), cpu_count()),
111         )
112 
113         dataset = dataset.filter(lambda data: data["labels"] is not None)
114         dataset = dataset.sort(column_names=("seq_category", "seq_length"), reverse=False, keep_in_memory=False)
115 
116         # We don't concatenate data samples here.
117         spliced_dataset = dataset
118         # Save each jsonl spliced dataset.
119         output_index = "0" * (5 - len(str(index))) + str(index)
120         output_name = f"part-{output_index}"
121         output_jsonl_path = os.path.join(args.data_jsonl_output_dir, output_name + ".jsonl")
122         # st = time.time()
123         with open(file=output_jsonl_path, mode="w", encoding="utf-8") as fp_writer:
124             spliced_count = 0
125             for spliced_data_point in spliced_dataset:
126                 if spliced_count % 500 == 0:
127                     logger.info(f"processing {spliced_count} spliced data points for {fp_writer.name}")
128                 spliced_count += 1
129                 fp_writer.write(json.dumps(spliced_data_point, ensure_ascii=False) + "\n")
130 
131         # Save each arrow spliced dataset
132         output_arrow_path = os.path.join(args.data_arrow_output_dir, output_name)
133         logger.info(f"Start to save {output_arrow_path}")
134         spliced_dataset = load_dataset(
135             path="json",
136             data_files=[output_jsonl_path],
137             cache_dir=os.path.join(args.data_cache_dir, "spliced_and_tokenized"),
138             keep_in_memory=False,
139             num_proc=cpu_count(),
140             split="train",
141         )
142         spliced_dataset.save_to_disk(dataset_path=output_arrow_path, num_proc=min(len(spliced_dataset), cpu_count()))
143 
144 
145 if __name__ == "__main__":
146     main()
147
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/applications/Colossal-LLaMA/prepare_sft_dataset.py b/applications/Colossal-LLaMA/prepare_sft_dataset.py
--- a/applications/Colossal-LLaMA/prepare_sft_dataset.py
+++ b/applications/Colossal-LLaMA/prepare_sft_dataset.py
@@ -10,7 +10,7 @@
 import os
 from multiprocessing import cpu_count
 
-from colossal_llama.dataset.conversation import default_conversation
+from colossal_llama.dataset.conversation import LLaMA2_Conv
 from colossal_llama.dataset.spliced_and_tokenized_dataset import supervised_tokenize_sft
 from datasets import dataset_dict, load_dataset
 from transformers import AddedToken, AutoTokenizer
@@ -78,6 +78,7 @@
     # Fix </s> split issue: https://github.com/huggingface/transformers/issues/23833
     if args.llama_version == 2:
         tokenizer.add_tokens(AddedToken("</s>", normalized=False, special=True), special_tokens=True)
+        default_conversation = LLaMA2_Conv
 
     tokenizer.add_bos_token = False
     tokenizer.add_eos_token = False

{"golden_diff": "diff --git a/applications/Colossal-LLaMA/prepare_sft_dataset.py b/applications/Colossal-LLaMA/prepare_sft_dataset.py\n--- a/applications/Colossal-LLaMA/prepare_sft_dataset.py\n+++ b/applications/Colossal-LLaMA/prepare_sft_dataset.py\n@@ -10,7 +10,7 @@\n import os\n from multiprocessing import cpu_count\n \n-from colossal_llama.dataset.conversation import default_conversation\n+from colossal_llama.dataset.conversation import LLaMA2_Conv\n from colossal_llama.dataset.spliced_and_tokenized_dataset import supervised_tokenize_sft\n from datasets import dataset_dict, load_dataset\n from transformers import AddedToken, AutoTokenizer\n@@ -78,6 +78,7 @@\n     # Fix </s> split issue: https://github.com/huggingface/transformers/issues/23833\n     if args.llama_version == 2:\n         tokenizer.add_tokens(AddedToken(\"</s>\", normalized=False, special=True), special_tokens=True)\n+        default_conversation = LLaMA2_Conv\n \n     tokenizer.add_bos_token = False\n     tokenizer.add_eos_token = False\n", "issue": "[BUG]: fine train llama-2-7b-hf prepare data set error , `bos_token` and `eos_token` should be the same with `conversation_template.seps`.\n### Is there an existing issue for this bug?\n\n- [X] I have searched the existing issues\n\n### \ud83d\udc1b Describe the bug\n\n python prepare_sft_dataset.py --data_input_dirs /models/train-data-dir  --tokenizer_dir /models/llama-2-13b-hf --data_output_dirs /models/train-data-dir/out  --num_spliced_dataset_bins 1\r\n\r\n\r\n\r\n\r\n File \"/ColossalAI/applications/Colossal-LLaMA/colossal_llama/dataset/spliced_and_tokenized_dataset.py\", line 90, in supervised_tokenize_sft\r\n    tokenizer.bos_token == conversation_template.seps[0] and tokenizer.eos_token == conversation_template.seps[1]\r\nAssertionError: `bos_token` and `eos_token` should be the same with `conversation_template.seps`.\r\n\"\"\"\r\n\r\nThe above exception was the direct cause of the following exception:\r\n\r\nTraceback (most recent call last):\r\n  File \"/ColossalAI/applications/Colossal-LLaMA/prepare_sft_dataset.py\", line 146, in <module>\r\n    main()\r\n  File \"/ColossalAI/applications/Colossal-LLaMA/prepare_sft_dataset.py\", line 102, in main\r\n    dataset = dataset.map(\r\n  File \"/usr/local/lib/python3.10/site-packages/datasets/arrow_dataset.py\", line 602, in wrapper\r\n    out: Union[\"Dataset\", \"DatasetDict\"] = func(self, *args, **kwargs)\r\n  File \"/usr/local/lib/python3.10/site-packages/datasets/arrow_dataset.py\", line 567, in wrapper\r\n    out: Union[\"Dataset\", \"DatasetDict\"] = func(self, *args, **kwargs)\r\n  File \"/usr/local/lib/python3.10/site-packages/datasets/arrow_dataset.py\", line 3248, in map\r\n    for rank, done, content in iflatmap_unordered(\r\n  File \"/usr/local/lib/python3.10/site-packages/datasets/utils/py_utils.py\", line 718, in iflatmap_unordered\r\n    [async_result.get(timeout=0.05) for async_result in async_results]\r\n  File \"/usr/local/lib/python3.10/site-packages/datasets/utils/py_utils.py\", line 718, in <listcomp>\r\n    [async_result.get(timeout=0.05) for async_result in async_results]\r\n  File \"/usr/local/lib/python3.10/site-packages/multiprocess/pool.py\", line 774, in get\r\n    raise self._value\r\n**AssertionError: `bos_token` and `eos_token` should be the same with `conversation_template.seps`.**\n\n### Environment\n\n------------ Environment ------------\r\nColossal-AI version: 0.3.7\r\nPyTorch version: 2.1.0\r\nSystem CUDA version: 11.0\r\nCUDA version required by PyTorch: 11.8\r\n\r\ntokenizer_config.json \r\n\r\n  \"add_bos_token\": true,\r\n  \"add_eos_token\": false,\r\n  \"bos_token\": {\r\n    \"__type\": \"AddedToken\",\r\n    \"content\": \"<s>\",\r\n    \"lstrip\": false,\r\n    \"normalized\": true,\r\n    \"rstrip\": false,\r\n    \"single_word\": false\r\n  },\r\n  \"clean_up_tokenization_spaces\": false,\r\n  \"eos_token\": {\r\n    \"__type\": \"AddedToken\",\r\n    \"content\": \"</s>\",\r\n    \"lstrip\": false,\r\n    \"normalized\": true,\r\n    \"rstrip\": false,\r\n    \"single_word\": false\r\n  },\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nPrepare sft dataset for fine-tuning\n\"\"\"\n\nimport argparse\nimport json\nimport math\nimport os\nfrom multiprocessing import cpu_count\n\nfrom colossal_llama.dataset.conversation import default_conversation\nfrom colossal_llama.dataset.spliced_and_tokenized_dataset import supervised_tokenize_sft\nfrom datasets import dataset_dict, load_dataset\nfrom transformers import AddedToken, AutoTokenizer\n\nfrom colossalai.logging import get_dist_logger\n\nlogger = get_dist_logger()\n\n\ndef main():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\n        \"--data_input_dirs\",\n        type=str,\n        required=True,\n        default=None,\n        help=\"Comma(i.e., ',') separated list of all data directories containing `.jsonl` data files.\",\n    )\n    parser.add_argument(\n        \"--tokenizer_dir\", type=str, required=True, default=None, help=\"A directory containing the tokenizer\"\n    )\n    parser.add_argument(\"--data_output_dirs\", type=str, default=\"data_output_dirs\", help=\"Data output directory\")\n    parser.add_argument(\"--max_length\", type=int, default=8192, help=\"Max length of each spliced tokenized sequence\")\n    parser.add_argument(\"--num_spliced_dataset_bins\", type=int, default=10, help=\"Number of spliced dataset bins\")\n    parser.add_argument(\"--llama_version\", type=int, default=3, help=\"LLaMA version\")\n    args = parser.parse_args()\n\n    if args.num_spliced_dataset_bins >= 100000:\n        raise ValueError(\"Too many spliced divisions, must be smaller than 100000\")\n\n    args.data_cache_dir = os.path.join(args.data_output_dirs, \"cache\")\n    args.data_jsonl_output_dir = os.path.join(args.data_output_dirs, \"jsonl\")\n    args.data_arrow_output_dir = os.path.join(args.data_output_dirs, \"arrow\")\n\n    if not os.path.exists(args.data_cache_dir):\n        os.makedirs(args.data_cache_dir)\n    if not os.path.exists(args.data_jsonl_output_dir):\n        os.makedirs(args.data_jsonl_output_dir)\n    if not os.path.exists(args.data_arrow_output_dir):\n        os.makedirs(args.data_arrow_output_dir)\n\n    # Prepare to all input datasets\n    input_data_paths = []\n    input_data_dirs = args.data_input_dirs.split(\",\")\n    for ds_dir in input_data_dirs:\n        ds_dir = os.path.abspath(ds_dir)\n        assert os.path.exists(ds_dir), f\"Not find data dir {ds_dir}\"\n        ds_files = [name for name in os.listdir(ds_dir) if name.endswith(\".jsonl\")]\n        ds_paths = [os.path.join(ds_dir, name) for name in ds_files]\n        input_data_paths.extend(ds_paths)\n\n    # Prepare to data splitting.\n    train_splits = []\n    split_interval = math.ceil(100 / args.num_spliced_dataset_bins)\n    for i in range(0, 100, split_interval):\n        start = i\n        end = i + split_interval\n        if end > 100:\n            end = 100\n        train_splits.append(f\"train[{start}%:{end}%]\")\n\n    # Prepare to the tokenizer.\n    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir)\n\n    # Fix </s> split issue: https://github.com/huggingface/transformers/issues/23833\n    if args.llama_version == 2:\n        tokenizer.add_tokens(AddedToken(\"</s>\", normalized=False, special=True), special_tokens=True)\n\n    tokenizer.add_bos_token = False\n    tokenizer.add_eos_token = False\n    if tokenizer.pad_token is None:\n        if tokenizer.unk_token is not None:\n            tokenizer.pad_token = tokenizer.unk_token\n        else:\n            tokenizer.pad_token = tokenizer.eos_token\n            tokenizer.unk_token = tokenizer.eos_token\n\n    list_dataset = load_dataset(\n        path=\"json\",\n        data_files=input_data_paths,\n        cache_dir=os.path.join(args.data_cache_dir, \"raw\"),\n        keep_in_memory=False,\n        split=train_splits,\n        num_proc=cpu_count(),\n    )\n    for index, dataset in enumerate(list_dataset):\n        assert isinstance(dataset, dataset_dict.Dataset)\n        logger.info(f\"Start to process part-{index}/{len(list_dataset)} of all original datasets.\")\n        dataset = dataset.map(\n            function=supervised_tokenize_sft,\n            fn_kwargs={\n                \"tokenizer\": tokenizer,\n                \"conversation_template\": default_conversation,\n                \"max_length\": args.max_length,\n            },\n            keep_in_memory=False,\n            num_proc=min(len(dataset), cpu_count()),\n        )\n\n        dataset = dataset.filter(lambda data: data[\"labels\"] is not None)\n        dataset = dataset.sort(column_names=(\"seq_category\", \"seq_length\"), reverse=False, keep_in_memory=False)\n\n        # We don't concatenate data samples here.\n        spliced_dataset = dataset\n        # Save each jsonl spliced dataset.\n        output_index = \"0\" * (5 - len(str(index))) + str(index)\n        output_name = f\"part-{output_index}\"\n        output_jsonl_path = os.path.join(args.data_jsonl_output_dir, output_name + \".jsonl\")\n        # st = time.time()\n        with open(file=output_jsonl_path, mode=\"w\", encoding=\"utf-8\") as fp_writer:\n            spliced_count = 0\n            for spliced_data_point in spliced_dataset:\n                if spliced_count % 500 == 0:\n                    logger.info(f\"processing {spliced_count} spliced data points for {fp_writer.name}\")\n                spliced_count += 1\n                fp_writer.write(json.dumps(spliced_data_point, ensure_ascii=False) + \"\\n\")\n\n        # Save each arrow spliced dataset\n        output_arrow_path = os.path.join(args.data_arrow_output_dir, output_name)\n        logger.info(f\"Start to save {output_arrow_path}\")\n        spliced_dataset = load_dataset(\n            path=\"json\",\n            data_files=[output_jsonl_path],\n            cache_dir=os.path.join(args.data_cache_dir, \"spliced_and_tokenized\"),\n            keep_in_memory=False,\n            num_proc=cpu_count(),\n            split=\"train\",\n        )\n        spliced_dataset.save_to_disk(dataset_path=output_arrow_path, num_proc=min(len(spliced_dataset), cpu_count()))\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "applications/Colossal-LLaMA/prepare_sft_dataset.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nPrepare sft dataset for fine-tuning\n\"\"\"\n\nimport argparse\nimport json\nimport math\nimport os\nfrom multiprocessing import cpu_count\n\nfrom colossal_llama.dataset.conversation import LLaMA2_Conv\nfrom colossal_llama.dataset.spliced_and_tokenized_dataset import supervised_tokenize_sft\nfrom datasets import dataset_dict, load_dataset\nfrom transformers import AddedToken, AutoTokenizer\n\nfrom colossalai.logging import get_dist_logger\n\nlogger = get_dist_logger()\n\n\ndef main():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\n        \"--data_input_dirs\",\n        type=str,\n        required=True,\n        default=None,\n        help=\"Comma(i.e., ',') separated list of all data directories containing `.jsonl` data files.\",\n    )\n    parser.add_argument(\n        \"--tokenizer_dir\", type=str, required=True, default=None, help=\"A directory containing the tokenizer\"\n    )\n    parser.add_argument(\"--data_output_dirs\", type=str, default=\"data_output_dirs\", help=\"Data output directory\")\n    parser.add_argument(\"--max_length\", type=int, default=8192, help=\"Max length of each spliced tokenized sequence\")\n    parser.add_argument(\"--num_spliced_dataset_bins\", type=int, default=10, help=\"Number of spliced dataset bins\")\n    parser.add_argument(\"--llama_version\", type=int, default=3, help=\"LLaMA version\")\n    args = parser.parse_args()\n\n    if args.num_spliced_dataset_bins >= 100000:\n        raise ValueError(\"Too many spliced divisions, must be smaller than 100000\")\n\n    args.data_cache_dir = os.path.join(args.data_output_dirs, \"cache\")\n    args.data_jsonl_output_dir = os.path.join(args.data_output_dirs, \"jsonl\")\n    args.data_arrow_output_dir = os.path.join(args.data_output_dirs, \"arrow\")\n\n    if not os.path.exists(args.data_cache_dir):\n        os.makedirs(args.data_cache_dir)\n    if not os.path.exists(args.data_jsonl_output_dir):\n        os.makedirs(args.data_jsonl_output_dir)\n    if not os.path.exists(args.data_arrow_output_dir):\n        os.makedirs(args.data_arrow_output_dir)\n\n    # Prepare to all input datasets\n    input_data_paths = []\n    input_data_dirs = args.data_input_dirs.split(\",\")\n    for ds_dir in input_data_dirs:\n        ds_dir = os.path.abspath(ds_dir)\n        assert os.path.exists(ds_dir), f\"Not find data dir {ds_dir}\"\n        ds_files = [name for name in os.listdir(ds_dir) if name.endswith(\".jsonl\")]\n        ds_paths = [os.path.join(ds_dir, name) for name in ds_files]\n        input_data_paths.extend(ds_paths)\n\n    # Prepare to data splitting.\n    train_splits = []\n    split_interval = math.ceil(100 / args.num_spliced_dataset_bins)\n    for i in range(0, 100, split_interval):\n        start = i\n        end = i + split_interval\n        if end > 100:\n            end = 100\n        train_splits.append(f\"train[{start}%:{end}%]\")\n\n    # Prepare to the tokenizer.\n    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir)\n\n    # Fix </s> split issue: https://github.com/huggingface/transformers/issues/23833\n    if args.llama_version == 2:\n        tokenizer.add_tokens(AddedToken(\"</s>\", normalized=False, special=True), special_tokens=True)\n        default_conversation = LLaMA2_Conv\n\n    tokenizer.add_bos_token = False\n    tokenizer.add_eos_token = False\n    if tokenizer.pad_token is None:\n        if tokenizer.unk_token is not None:\n            tokenizer.pad_token = tokenizer.unk_token\n        else:\n            tokenizer.pad_token = tokenizer.eos_token\n            tokenizer.unk_token = tokenizer.eos_token\n\n    list_dataset = load_dataset(\n        path=\"json\",\n        data_files=input_data_paths,\n        cache_dir=os.path.join(args.data_cache_dir, \"raw\"),\n        keep_in_memory=False,\n        split=train_splits,\n        num_proc=cpu_count(),\n    )\n    for index, dataset in enumerate(list_dataset):\n        assert isinstance(dataset, dataset_dict.Dataset)\n        logger.info(f\"Start to process part-{index}/{len(list_dataset)} of all original datasets.\")\n        dataset = dataset.map(\n            function=supervised_tokenize_sft,\n            fn_kwargs={\n                \"tokenizer\": tokenizer,\n                \"conversation_template\": default_conversation,\n                \"max_length\": args.max_length,\n            },\n            keep_in_memory=False,\n            num_proc=min(len(dataset), cpu_count()),\n        )\n\n        dataset = dataset.filter(lambda data: data[\"labels\"] is not None)\n        dataset = dataset.sort(column_names=(\"seq_category\", \"seq_length\"), reverse=False, keep_in_memory=False)\n\n        # We don't concatenate data samples here.\n        spliced_dataset = dataset\n        # Save each jsonl spliced dataset.\n        output_index = \"0\" * (5 - len(str(index))) + str(index)\n        output_name = f\"part-{output_index}\"\n        output_jsonl_path = os.path.join(args.data_jsonl_output_dir, output_name + \".jsonl\")\n        # st = time.time()\n        with open(file=output_jsonl_path, mode=\"w\", encoding=\"utf-8\") as fp_writer:\n            spliced_count = 0\n            for spliced_data_point in spliced_dataset:\n                if spliced_count % 500 == 0:\n                    logger.info(f\"processing {spliced_count} spliced data points for {fp_writer.name}\")\n                spliced_count += 1\n                fp_writer.write(json.dumps(spliced_data_point, ensure_ascii=False) + \"\\n\")\n\n        # Save each arrow spliced dataset\n        output_arrow_path = os.path.join(args.data_arrow_output_dir, output_name)\n        logger.info(f\"Start to save {output_arrow_path}\")\n        spliced_dataset = load_dataset(\n            path=\"json\",\n            data_files=[output_jsonl_path],\n            cache_dir=os.path.join(args.data_cache_dir, \"spliced_and_tokenized\"),\n            keep_in_memory=False,\n            num_proc=cpu_count(),\n            split=\"train\",\n        )\n        spliced_dataset.save_to_disk(dataset_path=output_arrow_path, num_proc=min(len(spliced_dataset), cpu_count()))\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "applications/Colossal-LLaMA/prepare_sft_dataset.py"}]}