scripts/train_tokenizer_32k.py

#!/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()