File size: 2,413 Bytes
a7532fa | 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 | import os
import tempfile
from datasets import load_dataset
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.processors import TemplateProcessing
BASE_DIR = r"E:\Aoban Selenia\Lightning"
PARQUET_FILE = os.path.join(
BASE_DIR,
"train0001of0001.parquet"
)
TOKENIZER_FILE = os.path.join(
BASE_DIR,
"lightning_tokenizer.json"
)
print(f"Loading dataset from {PARQUET_FILE}...")
dataset = load_dataset(
"parquet",
data_files=PARQUET_FILE,
split="train"
)
print(f"Loaded {len(dataset)} rows.")
temp_file = tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
suffix=".txt",
delete=False
)
temp_path = temp_file.name
print(f"Creating temporary tokenizer dataset...")
try:
with temp_file:
for row in dataset:
text = row["summary"]
if text and text.strip():
temp_file.write(text)
temp_file.write("\n")
print(f"Temporary dataset created at:")
print(temp_path)
tokenizer = Tokenizer(
BPE(
unk_token="<UNK>"
)
)
tokenizer.pre_tokenizer = ByteLevel(
add_prefix_space=True
)
tokenizer.decoder = ByteLevelDecoder()
trainer = BpeTrainer(
vocab_size=50000,
min_frequency=2,
special_tokens=[
"<PAD>",
"<UNK>",
"<|endoftext|>"
]
)
print("Training BPE tokenizer...")
tokenizer.train(
files=[temp_path],
trainer=trainer
)
tokenizer.post_processor = TemplateProcessing(
single="$A <|endoftext|>",
special_tokens=[
(
"<|endoftext|>",
tokenizer.token_to_id("<|endoftext|>")
)
]
)
tokenizer.save(
TOKENIZER_FILE
)
print()
print("Tokenizer trained successfully.")
print(f"Vocab Size: {tokenizer.get_vocab_size()}")
print(f"Saved to: {TOKENIZER_FILE}")
finally:
if os.path.exists(temp_path):
os.remove(temp_path)
print("Temporary tokenizer dataset deleted.")
|