Update README.md

README.md

@@ -452,60 +452,116 @@ from datasets import load_dataset
 from tqdm import tqdm
 
 OUTPUT_DIR = "data/alpaca_cleaned_mixed"
-OUTPUT_FILE = os.path.join(OUTPUT_DIR, "train.bin")
-enc = tiktoken.get_encoding("gpt2")
-
-def format_prompt(instruction, input_text, output):
-    sys_msg = "### System:\nYou are SmaLLMPro, a helpful AI Assistant developed by LH-Tech AI.\n\n"
-    if input_text and input_text.strip() != "":
-        return f"{sys_msg}### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n{output}<|endoftext|>"
+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:
-        return f"{sys_msg}### Instruction:\n{instruction}\n\n### Response:\n{output}<|endoftext|>"
+        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():
-    print("🚀 Starting data preparation fo SmaLLMPro...")
+    np.random.seed(SEED)
+    print(f"🚀 Starting Prepare-Script for SmaLLMPro (350M Instruct)...")
+    print(f"📚 Tokenizer: {TOKENIZER_NAME}")
     
-    print("📥 Loading Alpaca-Cleaned dataset from Huggingface...")
-    alpaca = load_dataset("yahma/alpaca-cleaned", split='train')
-
-    identity_prompts = [
-        {"instruction": "Who are you?", "input": "", "output": "I am SmaLLMPro, a helpful AI Assistant developed by LH-Tech AI. What do you want to talk about?"},
-        {"instruction": "Who developed you?", "input": "", "output": "I was developed by LH-Tech AI. Is there anything I can assist you with?"},
-        {"instruction": "What is you name?", "input": "", "output": "My name is SmaLLMPro. How can I help you today?"}
-    ]
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-    print("📥 Loading small FineWeb-Edu subset (anti-forgetting)...")
+    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)
-    fw_iter = iter(fineweb)
-
+    
     all_tokens = []
-
-    print("📝 Adding identity...")
-    for p in identity_prompts:
-        text = format_prompt(p['instruction'], p['input'], p['output'])
-        all_tokens.extend(enc.encode(text, allowed_special={"<|endoftext|>"}))
-
-    print("📝 Adding Alpaca-Cleaned...")
+    all_masks = []
+    
+    print("⚙️  Processing Alpaca...")
     for ex in tqdm(alpaca, desc="Alpaca"):
-        text = format_prompt(ex['instruction'], ex['input'], ex['output'])
-        all_tokens.extend(enc.encode(text, allowed_special={"<|endoftext|>"}))
+        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("📝 Adding FineWeb-Edu (Anti-Forgetting)...")
-    for _ in tqdm(range(2500), desc="FineWeb"):
+    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'] + "<|endoftext|>"
-            all_tokens.extend(enc.encode(text, allowed_special={"<|endoftext|>"}))
+            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)")
 
-    print(f"💾 Converting to numpy (Tokens: {len(all_tokens):,})...")
-    arr = np.array(all_tokens, dtype=np.uint16)
+    total_tokens = len(all_tokens)
+    print(f"\n💾 Saving {total_tokens:,} Tokens in '{OUTPUT_DIR}'...")
     
-    os.makedirs(OUTPUT_DIR, exist_ok=True)
-    arr.tofile(OUTPUT_FILE)
+    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.")
     
-    print(f"\n✅ Done! Binary file saved as: {OUTPUT_FILE}")
+    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()
@@ -521,24 +577,24 @@ from model import GPTConfig, GPT
 
 import numpy as np
 
-out_dir = '/home/user/checkpoints/350m_SmaLLMPro_Final'
+out_dir = '/home/user/350m_SmaLLMPro_Final'
 init_from = '/home/user/350m_fineweb'
 dataset = 'alpaca_cleaned_mixed'
 
 batch_size = 4
 gradient_accumulation_steps = 32
 block_size = 1024
-learning_rate = 3e-5
-max_iters = 3000
+learning_rate = 2e-5
+max_iters = 1500
 weight_decay = 0.1
 dropout = 0.1
-warmup_iters = 100
+warmup_iters = 0
 min_lr = 3e-6
 beta1, beta2 = 0.9, 0.95
 device = 'cuda'
 dtype = 'bfloat16'
 compile = True
-save_interval = 1000
+save_interval = 500
 
 os.makedirs(out_dir, exist_ok=True)
 torch.manual_seed(1337)
@@ -548,18 +604,23 @@ ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
 
 data_dir = os.path.join('data', dataset)
 train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
+train_mask = np.memmap(os.path.join(data_dir, 'train_mask.bin'), dtype=np.uint8, mode='r')
 
 def get_batch():
     ix = torch.randint(len(train_data) - block_size, (batch_size,))
     x = torch.stack([torch.from_numpy((train_data[i:i+block_size]).astype(np.int64)) for i in ix])
     y = torch.stack([torch.from_numpy((train_data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
+    m = torch.stack([torch.from_numpy((train_mask[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
+    
+    y[m == 0] = -100
+    
     x, y = x.to(device), y.to(device)
     return x, y
 
 print(f"📥 Loading Pretraining-Checkpoint from {init_from}...")
 ckpt_files = sorted([f for f in os.listdir(init_from) if f.endswith('.pt')])
 if not ckpt_files:
-    raise FileNotFoundError("No checkpoint in init_from folder found!")
+    raise FileNotFoundError("No checkpoint found in init_from directory!")
 
 ckpt_path = os.path.join(init_from, ckpt_files[-1])
 checkpoint = torch.load(ckpt_path, map_location=device)
@@ -576,7 +637,7 @@ model.load_state_dict(state_dict)
 model.to(device)
 
 if compile:
-    print("🚀 Compiling model...")
+    print("🚀 Compiling Model...")
     model = torch.compile(model)
 
 optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
@@ -589,7 +650,7 @@ def get_lr(it):
     coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
     return min_lr + coeff * (learning_rate - min_lr)
 
-print(f"🛠️ Starting finetuning...")
+print(f"🛠️ Starting Finetuning...")
 model.train()
 t0 = time.time()
 
@@ -613,7 +674,7 @@ for iter_num in range(max_iters + 1):
 
     if iter_num % 10 == 0:
         dt = time.time() - t0
-        print(f"Iter {iter_num}: Loss {loss.item()*gradient_accumulation_steps:.4f}, Zeit {dt*1000:.2f}ms, LR {lr:.2e}")
+        print(f"Iter {iter_num}: Loss {loss.item()*gradient_accumulation_steps:.4f}, Time {dt*1000:.2f}ms, LR {lr:.2e}")
         t0 = time.time()
 
     if iter_num > 0 and iter_num % save_interval == 0:
@@ -629,7 +690,7 @@ for iter_num in range(max_iters + 1):
         }
         torch.save(checkpoint_data, save_path)
 
-print(f"💾 Finetuning done. Saving SmaLLMPro final checkpoint...")
+print(f"💾 Finetuning done. Saving SmaLLMPro...")
 final_checkpoint = {
     'model': model.state_dict() if not compile else model._orig_mod.state_dict(),
     'model_args': checkpoint['model_args'],
@@ -646,41 +707,45 @@ import torch
 import tiktoken
 from model import GPTConfig, GPT
 
-# --- Config ---
-ckpt_path = '/home/user/350m_SmaLLMPro_Final/SmaLLMPro_iter_3000.pt'
+ckpt_path = '/home/user/350m_SmaLLMPro_Final/SmaLLMPro_iter_1500.pt'
 device = 'cuda'
 enc = tiktoken.get_encoding("gpt2")
 
-print("Loading SmaLLMPro...")
+print("Loading SmaLLMPro model checkpoint...")
 checkpoint = torch.load(ckpt_path, map_location=device)
 gptconf = GPTConfig(**checkpoint['model_args'])
 model = GPT(gptconf)
-model.load_state_dict(checkpoint['model'])
+
+state_dict = checkpoint['model']
+unwanted_prefix = '_orig_mod.'
+for k,v in list(state_dict.items()):
+    if k.startswith(unwanted_prefix):
+        state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+
+model.load_state_dict(state_dict)
 model.eval()
 model.to(device)
-print("Ready!\n")
+print(f"Model {ckpt_path} ready!\n")
 
 def run_chat():
-    print("--- SmaLLMPro Chatbot (Type 'exit' to quit) ---")
-    
-    sys_msg = "### System:\nYou are SmaLLMPro, a helpful AI Assistant developed by LH-Tech AI.\n\n"
+    print("--- SmaLLMPro Chatbot (Type 'exit' to leave) ---")
     
     while True:
         user_input = input("You: ")
-        if user_input.lower() in ["exit", "quit", "beenden"]:
+        if user_input.lower() in ["exit", "quit"]:
             break
 
-        prompt = f"{sys_msg}### Instruction:\n{user_input}\n\n### Response:\n"
+        prompt = f"Instruction:\n{user_input}\n\nResponse:\n"
         
         x = torch.tensor(enc.encode(prompt), dtype=torch.long, device=device)[None, ...]
         
         print("SmaLLMPro: ", end="", flush=True)
         with torch.no_grad():
             with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
-                y = model.generate(x, max_new_tokens=256, temperature=0.7, top_k=40)
+                y = model.generate(x, max_New_tokens=256, temperature=0.4, top_k=25)
                 full_text = enc.decode(y[0].tolist())
                 
-                response = full_text.split("### Response:\n")[-1].split("<|endoftext|>")[0].strip()
+                response = full_text.split("Response:\n")[-1].split("<|endoftext|>")[0].strip()
                 print(response + "\n")
 
 if __name__ == "__main__":