-
-
Notifications
You must be signed in to change notification settings - Fork 837
/
make_data.py
159 lines (137 loc) · 5.13 KB
/
make_data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
import json, math, random, sys, time, shutil, os, string, re, fileinput
import numpy as np
"""
How to use:
python make_data.py demo.jsonl 3 4096
This will:
==> shuffle & duplicate demo.jsonl (for 3 epochs, good for finetuning) note: this will be very slow for large jsonl and we need more efficient code.
==> load jsonl and tokenize
==> save as demo.bin & demo.idx
==> compute "magic_prime" for ctxlen 4096
Example:
Assume your source jsonl is:
{"text":"aa"}
{"text":"bb"}
{"text":"cc"}
{"text":"dd"}
The final binidx will be like (here "/" means end_of_doc, which is actually token [0]):
bb/aa/dd/cc/dd/aa/bb/cc/dd/bb/cc/aa/
where the data is repeated 3 times (each time with different shuffle)
"""
########################################################################################################
# MMapIndexedDatasetBuilder
########################################################################################################
from tokenizer.rwkv_tokenizer import TRIE_TOKENIZER
tokenizer = TRIE_TOKENIZER("tokenizer/rwkv_vocab_v20230424.txt")
from src.binidx import MMapIndexedDataset
def index_file_path(prefix_path):
return prefix_path + ".idx"
def data_file_path(prefix_path):
return prefix_path + ".bin"
class MMapIndexedDatasetBuilder(object):
def __init__(self, out_file, dtype=np.uint16):
self._data_file = open(out_file, "wb")
self._dtype = dtype
self._sizes = []
self._doc_idx = [0]
def add_item(self, np_array):
assert np_array.dtype == self._dtype
self._data_file.write(np_array.tobytes(order="C"))
self._sizes.append(np_array.size)
def end_document(self):
self._doc_idx.append(len(self._sizes))
def finalize(self, index_file):
self._data_file.close()
with MMapIndexedDataset.Index.writer(index_file, self._dtype) as index:
index.write(self._sizes, self._doc_idx)
cnt = 0
def add_raw(raw):
global builder, cnt
out = tokenizer.encode(raw)
if tokenizer.decode(out) != raw:
print("ERROR" * 100)
exit(0)
out.append(0) # [0] = end_of_doc for rwkv tokenizer
builder.add_item(np.array(out, dtype=np.uint16))
builder.end_document()
if cnt % 500 == 0:
print(cnt, end=" ", flush=True)
cnt += 1
def is_prime(n):
if n <= 1:
return False
if n <= 3:
return True
if n % 2 == 0 or n % 3 == 0:
return False
i = 5
while i * i <= n:
if n % i == 0 or n % (i + 2) == 0:
return False
i += 6
return True
########################################################################################################
N_EPOCH = int(sys.argv[2].strip())
IN_FILE = sys.argv[1].strip()
OUT_NAME = os.path.splitext(os.path.basename(IN_FILE))[0]
CTX_LEN = int(sys.argv[3].strip())
TEMP_FILE = "make_data_temp.jsonl"
print(f"### Convert {IN_FILE} to {OUT_NAME}.bin/idx...")
with open(IN_FILE, "r", encoding="utf-8") as file:
non_empty_lines = [line.strip() for line in file if line.strip()]
print(f"### Found {len(non_empty_lines)} non-empty lines in {IN_FILE}")
file = open(TEMP_FILE, "w", encoding="utf-8")
for i in range(N_EPOCH):
print(f"Shuffle: {i+1} out of {N_EPOCH}")
random.shuffle(non_empty_lines)
for entry in non_empty_lines:
file.write(entry + "\n")
file.close()
########################################################################################################
print("### Building binidx...")
builder = MMapIndexedDatasetBuilder(f"{OUT_NAME}.bin")
with fileinput.input(TEMP_FILE, encoding="utf-8") as ffff:
for line in ffff:
x = json.loads(line)["text"]
add_raw(x)
builder.finalize((f"{OUT_NAME}.idx"))
print("done")
print("### Verifying result...")
data = MMapIndexedDataset(OUT_NAME)
data_len = len(data)
data_size = len(data._bin_buffer) // data._index._dtype_size
TODO = [0, data_len - 1]
PREVIEW_LIMIT = 100
for idx in TODO:
ptr, size = data._index[idx]
dix = data.get(idx=idx, offset=0, length=size).astype(int)
print("-" * 70 + f"[{OUT_NAME} idx {idx} sz {size}]")
assert dix[-1] == 0
dix = dix[:-1]
if len(dix) > PREVIEW_LIMIT:
try:
print(tokenizer.decode(dix[:PREVIEW_LIMIT]))
except:
try:
print(tokenizer.decode(dix[: PREVIEW_LIMIT + 1]))
except:
print(tokenizer.decode(dix[: PREVIEW_LIMIT + 2]))
print("· " * 30)
try: # avoid utf-8 bug
print(tokenizer.decode(dix[-PREVIEW_LIMIT:]))
except:
try:
print(tokenizer.decode(dix[-PREVIEW_LIMIT - 1 :]))
except:
print(tokenizer.decode(dix[-PREVIEW_LIMIT - 2 :]))
else:
print(tokenizer.decode(dix))
print(f"{'-'*80}\n### Final {OUT_NAME}.bin/idx has {data_size} tokens, {data_len} items. Dtype {data._index.dtype}")
if data_size >= CTX_LEN * 3:
n_chunk = int(data_size // CTX_LEN) - 1
for i in range(n_chunk, 0, -1):
if i % 3 == 2:
if is_prime(i):
print(f"\n### magic_prime = {i} (for ctxlen {CTX_LEN})")
print(f'\n--my_exit_tokens {data_size} --magic_prime {i} --ctx_len {CTX_LEN}\n')
exit(0)