File size: 4,116 Bytes
a47f627 | 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 | import os
import numpy as np
import tiktoken
from datasets import load_dataset
from tqdm import tqdm
OUTPUT_DIR = "data/alpaca_cleaned_mixed"
TOKENIZER_NAME = "gpt2"
SEED = 1337
FINEWEB_SAMPLES = 2500
enc = tiktoken.get_encoding(TOKENIZER_NAME)
EOS_TOKEN = "<|endoftext|>"
def format_prompt_with_mask(instruction, input_text, output):
"""
Formatiert den Prompt und erstellt die Loss-Maske.
Format:
Instruction: ...
Input: ... (optional)
Response: ... <|endoftext|>
"""
if input_text and input_text.strip():
prompt_text = f"Instruction:\n{instruction}\n\nInput:\n{input_text}\n\nResponse:\n"
else:
prompt_text = f"Instruction:\n{instruction}\n\nResponse:\n"
completion_text = f"{output}{EOS_TOKEN}"
prompt_ids = enc.encode(prompt_text, allowed_special={'<|endoftext|>'})
completion_ids = enc.encode(completion_text, allowed_special={'<|endoftext|>'})
full_ids = prompt_ids + completion_ids
mask = [0] * len(prompt_ids) + [1] * len(completion_ids)
return full_ids, mask
def main():
np.random.seed(SEED)
print(f"🚀 Starting Prepare-Script for SmaLLMPro (350M Instruct)...")
print(f"📚 Tokenizer: {TOKENIZER_NAME}")
os.makedirs(OUTPUT_DIR, exist_ok=True)
print("📥 Loading 'yahma/alpaca-cleaned' (Chat-Instructions)...")
alpaca = load_dataset("yahma/alpaca-cleaned", split='train')
print(f"📥 Loading 'HuggingFaceFW/fineweb-edu' (Sample-10BT) for {FINEWEB_SAMPLES} Samples...")
fineweb = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split='train', streaming=True)
all_tokens = []
all_masks = []
print("⚙️ Processing Alpaca...")
for ex in tqdm(alpaca, desc="Alpaca"):
ids, mask = format_prompt_with_mask(ex['instruction'], ex['input'], ex['output'])
all_tokens.extend(ids)
all_masks.extend(mask)
alpaca_len = len(all_tokens)
print(f" -> Alpaca Tokens: {alpaca_len:,}")
print("⚙️ Processing FineWeb (Anti-Forgetting)...")
fw_iter = iter(fineweb)
fw_count = 0
fw_tokens_count = 0
for _ in tqdm(range(FINEWEB_SAMPLES), desc="FineWeb"):
try:
ex = next(fw_iter)
text = ex['text'] + EOS_TOKEN
ids = enc.encode(text, allowed_special={EOS_TOKEN})
all_tokens.extend(ids)
all_masks.extend([1] * len(ids))
fw_tokens_count += len(ids)
fw_count += 1
except StopIteration:
break
print(f" -> FineWeb Tokens: {fw_tokens_count:,} (from {fw_count} documents)")
total_tokens = len(all_tokens)
print(f"\n💾 Saving {total_tokens:,} Tokens in '{OUTPUT_DIR}'...")
token_arr = np.array(all_tokens, dtype=np.uint16)
token_arr.tofile(os.path.join(OUTPUT_DIR, "train.bin"))
mask_arr = np.array(all_masks, dtype=np.uint8)
mask_arr.tofile(os.path.join(OUTPUT_DIR, "train_mask.bin"))
print("\n🔍 --- SANITY CHECK ---")
print("I decode the first 50 tokens of the first sample, to check, if everything is okay.")
print("Green (TRAIN) = The things the model learns. Grey (IGNORE) = The things the model only reads.")
check_len = 100
sample_ids = all_tokens[:check_len]
sample_mask = all_masks[:check_len]
decoded_parts = []
for t_id, m_val in zip(sample_ids, sample_mask):
token_str = enc.decode([t_id])
if m_val == 1:
decoded_parts.append(f"\033[92m{token_str}\033[0m")
else:
decoded_parts.append(f"\033[90m{token_str}\033[0m")
print("".join(decoded_parts))
print("\n(Legend: \033[90mGrey=Prompt/Ignored\033[0m, \033[Green=Response/Learned\033[0m)")
if len(token_arr) != len(mask_arr):
print("\n❌ Warning: Token and Mask Array have different lengths! Something has gone wrong!")
else:
print("\n✅ Everything seems to be fine. The arrays are synchronized. You can now start the training.")
if __name__ == "__main__":
main() |