#!/usr/bin/env python3 """ Train a 32K BPE tokenizer on text data from HuggingFace datasets. Falls back to dummy data if downloads fail. Saves to checkpoints/tokenizer_32k.json. Usage: python scripts/train_tokenizer_32k.py """ import sys import os from pathlib import Path # Add project root to path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from model.tokenizer import BPETokenizer def load_sharegpt_texts(max_samples: int = 5000) -> list[str]: """Load text from ShareGPT dataset.""" try: from datasets import load_dataset ds = load_dataset("anon8231489123/ShareGPT_Vicuna_unfiltered", split="train", streaming=True) texts = [] for i, example in enumerate(ds): if i >= max_samples: break conversations = example.get("conversations", []) for turn in conversations: val = turn.get("value", "") if val and len(val) > 20: texts.append(val) print(f" ShareGPT: loaded {len(texts)} text segments") return texts except Exception as e: print(f" ShareGPT: download failed ({e}), skipping") return [] def load_coco_captions(max_samples: int = 5000) -> list[str]: """Load COCO captions from HuggingFace.""" try: from datasets import load_dataset ds = load_dataset("HuggingFaceM4/COCO", split="train", streaming=True) texts = [] for i, example in enumerate(ds): if i >= max_samples: break caption = example.get("caption", "") or example.get("sentences", [""])[0] if isinstance(caption, list): caption = caption[0] if caption else "" if caption and len(caption) > 10: texts.append(caption) print(f" COCO captions: loaded {len(texts)} captions") return texts except Exception as e: print(f" COCO captions: download failed ({e}), skipping") return [] def load_vqa_answers(max_samples: int = 5000) -> list[str]: """Load VQAv2 questions and answers.""" try: from datasets import load_dataset ds = load_dataset("HuggingFaceM4/VQAv2", split="train", streaming=True) texts = [] for i, example in enumerate(ds): if i >= max_samples: break question = example.get("question", "") answers = example.get("answers", []) if question: texts.append(question) for ans in answers[:3]: a = ans if isinstance(ans, str) else ans.get("answer", "") if a: texts.append(a) print(f" VQAv2: loaded {len(texts)} text segments") return texts except Exception as e: print(f" VQAv2: download failed ({e}), skipping") return [] def generate_dummy_data() -> list[str]: """Generate dummy training data as fallback.""" print(" Using dummy fallback data") base_texts = [ "A person walking a dog in the park on a sunny day", "The camera detected motion near the entrance gate", "Count the number of vehicles in the parking lot", "Two people are standing near the reception desk", "Alert: unauthorized access detected at door 3", "The quick brown fox jumps over the lazy dog", "A red car is parked next to a blue truck", "Three workers wearing safety helmets on the construction site", "The temperature sensor reads 72 degrees Fahrenheit", "Multiple pedestrians crossing the street at the intersection", "Security camera footage shows normal activity in the lobby", "Object detection identified 5 cars and 2 motorcycles", "The license plate reads ABC 1234 on the white sedan", "A delivery truck is unloading packages at the loading dock", "Night vision mode activated due to low light conditions", "Face recognition matched employee ID 4521 at checkpoint", "Traffic flow analysis shows congestion on lane 2", "Smoke detected in zone B of the warehouse area", "The PTZ camera is tracking a moving object heading north", "Crowd density estimation shows approximately 150 people", ] # Repeat with variations to get enough training data texts = [] for i in range(500): for text in base_texts: texts.append(text) # Add variations texts.append(text.lower()) texts.append(text.upper()) words = text.split() if len(words) > 3: texts.append(" ".join(words[:len(words) // 2])) texts.append(" ".join(words[len(words) // 2:])) return texts def main(): print("=" * 60) print("Training 32K BPE Tokenizer") print("=" * 60) # Collect training data print("\nLoading training data...") all_texts = [] all_texts.extend(load_sharegpt_texts()) all_texts.extend(load_coco_captions()) all_texts.extend(load_vqa_answers()) # Fall back to dummy data if we got very little if len(all_texts) < 1000: print("\n Insufficient data from HuggingFace, adding dummy data...") all_texts.extend(generate_dummy_data()) print(f"\nTotal training texts: {len(all_texts):,}") # Train tokenizer print("\nTraining BPE tokenizer (vocab_size=32768)...") tok = BPETokenizer(vocab_size=32768) tok.train(all_texts) # Save save_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "checkpoints", "tokenizer_32k.json") Path(save_path).parent.mkdir(parents=True, exist_ok=True) tok.save(save_path) print(f"\nSaved tokenizer to: {save_path}") # Print stats print("\n" + "=" * 60) print("Vocab Statistics:") print("=" * 60) print(f" Total vocab size (configured): {tok.vocab_size:,}") print(f" Total tokens in vocab: {len(tok.vocab):,}") print(f" Special tokens: {len(tok.SPECIAL_TOKENS)}") print(f" Byte fallback tokens: {len(tok.BYTE_TOKENS)}") print(f" BPE merges learned: {len(tok.merges):,}") print(f" Character + merged tokens: {len(tok.vocab) - len(tok.SPECIAL_TOKENS) - len(tok.BYTE_TOKENS):,}") # Test encode/decode print("\n" + "=" * 60) print("Encode/Decode Examples:") print("=" * 60) test_texts = [ "A person walking in the park", "Count the vehicles in the lot", "Hello 你好 世界", "Alert: motion detected at zone B", ] for text in test_texts: ids = tok.encode(text) decoded = tok.decode(ids) print(f" Input: {text}") print(f" Tokens: {len(ids)} IDs") print(f" Decoded: {decoded}") print() if __name__ == "__main__": main()