File size: 1,872 Bytes
a8612eb | 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 | #!/usr/bin/env python3
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from mla.tokenizer import Tokenizer
TOK = Path("data/tokenizer/tokenizer.json")
SPLITS = {
"train": Path("data/sft/train.txt"),
"val": Path("data/sft/val.txt"),
}
OUT_DIR = Path("data/sft")
def build_mask(ids, bos, eos, user, assistant):
mask = []
role = None
for t in ids:
if t == bos:
mask.append(0)
elif t == user:
role = "u"
mask.append(0)
elif t == assistant:
role = "a"
mask.append(0)
else:
mask.append(1 if role == "a" else 0)
return mask
def main():
tok = Tokenizer.load(TOK)
bos = tok.special_to_id["<bos>"]
eos = tok.special_to_id["<eos>"]
user = tok.special_to_id["<|user|>"]
assistant = tok.special_to_id["<|assistant|>"]
OUT_DIR.mkdir(parents=True, exist_ok=True)
for name, path in SPLITS.items():
id_stream = []
mask_stream = []
n_dialogues = 0
for line in path.read_text(encoding="utf-8").splitlines():
if not line:
continue
ids = tok.encode(line)
mask = build_mask(ids, bos, eos, user, assistant)
id_stream.extend(ids)
mask_stream.extend(mask)
n_dialogues += 1
ids_arr = np.array(id_stream, dtype=np.uint16)
mask_arr = np.array(mask_stream, dtype=np.uint8)
np.save(OUT_DIR / f"{name}_ids.npy", ids_arr)
np.save(OUT_DIR / f"{name}_mask.npy", mask_arr)
frac = mask_arr.mean() if len(mask_arr) else 0.0
print(f"{name}: dialogues={n_dialogues} tokens={len(ids_arr)} "
f"assistant_frac={frac:.3f} -> {name}_ids.npy / {name}_mask.npy")
if __name__ == "__main__":
main()
|