Upload train_tokenizer.py

train_tokenizer.py

@@ -0,0 +1,172 @@
+# improved_train_tokenizer_v2.py
+
+import os
+import sys
+from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, processors, normalizers
+from transformers import PreTrainedTokenizerFast
+
+# --- Configuration ---
+TRAIN_FILES = ["improved_sentences.txt"]  # Use the preprocessed file
+VOCAB_SIZE = 32000
+SPECIAL_TOKENS = ["<pad>", "<unk>", "<s>", "</s>", "<mask>"]
+OUTPUT_DIR = "./improved_tokenizer_v2"
+
+# --- Input File Check ---
+if not TRAIN_FILES or not os.path.exists(TRAIN_FILES[0]):
+    print(f"Error: Training file '{TRAIN_FILES[0]}' not found.")
+    sys.exit(1)
+
+print(f"Starting tokenizer training...")
+print(f"Training file(s): {TRAIN_FILES}")
+print(f"Target vocab size: {VOCAB_SIZE}")
+print(f"Output directory: {OUTPUT_DIR}")
+
+# --- Initialize Tokenizer ---
+# We'll use ByteLevel BPE with proper whitespace handling
+tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
+
+# --- Set Normalizer ---
+# This helps standardize the text before tokenization
+tokenizer.normalizer = normalizers.Sequence([
+    normalizers.NFC(),  # Unicode normalization
+    normalizers.Replace(r"\s+", " ")  # Replace multiple spaces with a single space
+])
+
+# --- Set Pre-tokenizer ---
+# This is critical for handling whitespace correctly
+tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)  # Back to True for proper space handling
+print(f"Using pre-tokenizer: ByteLevel(add_prefix_space=True)")
+
+# --- Set Decoder ---
+tokenizer.decoder = decoders.ByteLevel()
+print(f"Using decoder: {tokenizer.decoder.__class__.__name__}")
+
+# --- Define Trainer ---
+trainer = trainers.BpeTrainer(
+    vocab_size=VOCAB_SIZE,
+    special_tokens=SPECIAL_TOKENS,
+    show_progress=True,
+    initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+)
+
+# --- Train Tokenizer ---
+print("\nTraining the tokenizer model (this might take a while)...")
+try:
+    tokenizer.train(files=TRAIN_FILES, trainer=trainer)
+    print("Training completed successfully.")
+except Exception as e:
+    print(f"\nError during tokenizer training: {e}")
+    sys.exit(1)
+
+# --- Add Post-processor ---
+tokenizer.post_processor = processors.TemplateProcessing(
+    single="<s> $A </s>",
+    pair="<s> $A </s> $B </s>",
+    special_tokens=[
+        ("<s>", tokenizer.token_to_id("<s>")),
+        ("</s>", tokenizer.token_to_id("</s>")),
+    ],
+)
+
+# --- Save Core Tokenizer ---
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+tokenizer_path = os.path.join(OUTPUT_DIR, "tokenizer.json")
+try:
+    tokenizer.save(tokenizer_path)
+    print(f"\nCore tokenizer saved to: {tokenizer_path}")
+except Exception as e:
+    print(f"Error saving core tokenizer: {e}")
+    sys.exit(1)
+
+# --- Create and Save HF Wrapper ---
+print("\nWrapping tokenizer with PreTrainedTokenizerFast...")
+try:
+    hf_tokenizer = PreTrainedTokenizerFast(
+        tokenizer_file=tokenizer_path,
+        unk_token="<unk>",
+        pad_token="<pad>",
+        cls_token="<s>",
+        sep_token="</s>",
+        mask_token="<mask>",
+        add_prefix_space=True  # Match the pre-tokenizer setting
+    )
+    hf_tokenizer.save_pretrained(OUTPUT_DIR)
+    print(f"Hugging Face compatible tokenizer files saved to: {OUTPUT_DIR}")
+except Exception as e:
+    print(f"Error saving Hugging Face tokenizer: {e}")
+    sys.exit(1)
+
+# --- Verification Step ---
+print("\n--- Verification ---")
+try:
+    print(f"Loading tokenizer for verification from: {OUTPUT_DIR}")
+    loaded_hf_tokenizer = PreTrainedTokenizerFast.from_pretrained(OUTPUT_DIR)
+
+    # Test multiple cases, especially those starting with periods or spaces
+    test_cases = [
+        "Simple sentence.",
+        " Sentence starting with space.",
+        "Sentence. Another sentence.",
+        ". Sentence starting with period.", 
+        "Word.Word",
+        "The quick brown fox jumps over the lazy dog."
+    ]
+    
+    print("\n=== Testing with new tokenizer ===")
+    for i, text in enumerate(test_cases):
+        print(f"\nTest {i+1}: '{text}'")
+        tokens = loaded_hf_tokenizer.tokenize(text)
+        print(f"Tokens: {tokens}")
+        
+        encoded = loaded_hf_tokenizer.encode(text, add_special_tokens=True)
+        decoded = loaded_hf_tokenizer.decode(encoded, skip_special_tokens=True)
+        print(f"Encoded: {encoded}")
+        print(f"Decoded: '{decoded}'")
+        
+        # Check if tokenization properly preserves content
+        if text.strip() == decoded.strip():
+            print("✓ Encoding/decoding preserved text content")
+        else:
+            print(f"⚠ Warning: Text content changed during encoding/decoding")
+            print(f"  Original: '{text}'")
+            print(f"  Decoded:  '{decoded}'")
+
+    # Check first token distributions
+    print("\n=== First Position Token Analysis ===")
+    print("Analyzing first token after <s> for potential bias...")
+    
+    # Simplified analysis of first token (just for demonstration)
+    from collections import Counter
+    first_token_counter = Counter()
+    
+    with open(TRAIN_FILES[0], 'r', encoding='utf-8') as f:
+        for i, line in enumerate(f):
+            if i >= 100:  # Just check first 100 lines
+                break
+            line = line.strip()
+            if not line:
+                continue
+                
+            encoded = loaded_hf_tokenizer.encode(line, add_special_tokens=True)
+            if len(encoded) > 1:  # Make sure there's at least one token after <s>
+                first_token_id = encoded[1]
+                first_token_counter[first_token_id] += 1
+    
+    total = sum(first_token_counter.values())
+    if total > 0:
+        print(f"\nTop 5 tokens at first position (after <s>) from {total} samples:")
+        for token_id, count in first_token_counter.most_common(5):
+            token_text = loaded_hf_tokenizer.decode([token_id])
+            percentage = (count / total) * 100
+            print(f"Token: '{token_text}' (ID: {token_id}) | Count: {count} | {percentage:.2f}%")
+            
+        # Specifically check period token
+        period_id = loaded_hf_tokenizer.encode('.', add_special_tokens=False)[0]
+        period_count = first_token_counter.get(period_id, 0)
+        period_percentage = (period_count / total) * 100 if total > 0 else 0
+        print(f"\nPeriod token ('.', ID: {period_id}) at first position: {period_count} times ({period_percentage:.2f}%)")
+    
+except Exception as e:
+    print(f"Error during verification: {e}")
+
+print("\n--- Tokenizer training script finished ---")