Create Surgical_Vocab_Reduction_for_QWEN2B_DotOCRv1.5_V1.0untested.py

Surgical_Vocab_Reduction_for_QWEN2B_DotOCRv1.5_V1.0untested.py

@@ -0,0 +1,482 @@
+print("starting Surgical_Vocab_Reduction_for_QWEN2B_DotOCRv1.5_V1.0untested.py")
+r"""
+surgical_vocab_reduction.py
+
+Description:
+    Surgically prunes the text vocabulary of Qwen2/dots.ocr-1.5 based models.
+    Removes Chinese, Cyrillic, Arabic, and other non-European characters from 
+    the tokenizer, while preserving English, French, Spanish, special vision 
+    tokens, layout coordinate symbols, and byte fallback sequences.
+    
+    Dynamically prunes model weight files (safetensors/bin) to adjust 
+    the input embeddings and output LM head dimensions to the new vocabulary size, 
+    updating config.json and generation configs.
+
+Dependencies:
+    - torch
+    - safetensors
+    - transformers
+
+Surgical Vocabulary & Parameter Reduction for Qwen2 / dots.ocr-1.5 (V1.0 Untested)
+
+This repository contains details and scripts for
+Surgical_Vocab_Reduction_for_QWEN2B_DotOCRv1.5_V1.0untested.py, a pipeline
+designed to prune the large multilingual vocabulary of the dots.ocr-1.5 model
+(and other Qwen2-based Vision-Language Models).
+
+By surgically removing non-target Unicode blocks (such as Chinese, Cyrillic,
+Arabic, Hebrew, Japanese, and Korean) while keeping English, French, Spanish,
+special layout coordinate symbols, and byte-fallback sequences, this script
+significantly reduces both the model's disk/VRAM footprint and its output logit
+calculation latency during local edge inference (such as on Intel Core Ultra
+NPUs or mobile devices).
+
+This implementation is inspired by and built upon the methodology outlined in
+Method for Dynamically Reducing Logit Computation in LLMs.
+
+1. The Core Problem
+
+Modern, state-of-the-art multimodal models like dots.ocr-1.5 are designed for
+global multilingual coverage. To achieve this, they utilize very large
+vocabularies (Qwen2 series, for example, features a vocabulary of 151,936
+tokens).
+
+While this makes the model highly capable globally, it introduces significant
+overhead for deployment scenarios limited to Western European scripts (e.g.,
+English, French, Spanish):
+
+1.  Parameter Inflation: The input embedding layer (model.embed_tokens.weight)
+    and the final Language Modeling head (lm_head.weight) scale linearly with
+    vocabulary size. At 151,936 tokens, these two layers alone consume hundreds
+    of millions of parameters, adding several gigabytes of unnecessary storage
+    and VRAM.
+2.  Computational Overhead (Logit Latency): During autoregressive token
+    generation, the final layer performs a matrix multiplication of shape
+    [batch_size * seq_len, hidden_dim] @ [hidden_dim, vocab_size] to compute
+    logits. Performing this massive projection on every decoded token introduces
+    latency, particularly on resource-constrained CPUs and edge NPUs.
+
+2. The Surgical Reduction Method
+
+The script executes a targeted, multi-stage compression process that shrinks the
+model's text processing layers without degrading its fundamental visual-spatial,
+OCR, or layout parsing intelligence:
+
+[ Original VLM ] ──> [ 1. Scan Tokenizer & Filter Vocab ] ──> [ 2. Filter BPE Merges ] 
+                           │ (ASCII, Latin-1, & Special Tags Only)       │
+                           ▼                                             ▼
+[ Pruned Weights ] <── [ 4. Prune Embedding/Head Tensors ] <── [ 3. Re-index IDs ]
+
+Stage A: Tokenizer Parsing & Robust Filtering
+
+The script reads tokenizer.json and evaluates each token string against strict,
+multi-tier criteria to determine which to keep:
+
+  - Preserve Special & Coordinate Tokens: Bounding box markers, segment
+    boundaries, and visual layout tags (e.g., <|image_1|>, <|box_start|>, etc.)
+    are detected and kept to protect the spatial parsing mechanism.
+  - Preserve Byte-Fallback Mappings: To prevent the tokenizer from failing when
+    encountering out-of-vocabulary characters, any token representing a single
+    byte (Latin-1 range 0x00 - 0xFF or length-1 strings) is retained.
+  - Targeted Script Deletion: Multi-character tokens containing Chinese
+    characters, Japanese Kana, Hangul, Cyrillic, Arabic, Devanagari, or Thai are
+    discarded.
+  - Targeted Language Retention: Standard ASCII punctuation, numbers, and the
+    accented character sets of English, French, and Spanish (using Latin-1 and
+    Latin Extended-A Unicode ranges) are preserved.
+
+Stage B: Re-Indexing & Merge List Sanitization
+
+  - Discarded tokens are stripped, and the remaining tokens are sequentially
+    re-indexed from 0 to N_{\text{new}} - 1.
+  - The Byte Pair Encoding (BPE) "merges" rules are filtered. If a merge pair
+    contains a sub-token or output token that was deleted, the merge rule is
+    safely omitted. This prevents tokenizer errors or lookup warnings.
+
+Stage C: Surgical Tensor Pruning
+
+  - The script scans the model directory for weights files (both PyTorch .bin
+    and Safetensors .safetensors formats).
+  - It locates the input embedding weights (model.embed_tokens.weight or
+    transformer.wte.weight) and the output projection head weights
+    (lm_head.weight).
+  - It extracts only the slices of these matrices corresponding to the indices
+    of the retained vocabulary, updating their shapes dynamically.
+
+Stage D: Configuration Re-alignment
+
+  - Updates config.json with the new "vocab_size".
+  - Remaps bos_token_id, eos_token_id, and pad_token_id to their new values
+    inside generation_config.json.
+
+3. Anticipated Benefits
+
+  - VRAM and Storage Compression: Shifting from 151,936 tokens down to an
+    estimated 30,000 \text{--} 40,000 tokens significantly shrinks the
+    parameters of the text layers. For a hidden size of D = 3584, this reduction
+    strips away over 800 million parameters, translating to immediate savings of
+    up to 1.6 GB of VRAM/storage in FP16 (or comparable savings under quantized
+    configurations).
+  - Improved Generation Latency: Computing output logits against a vocabulary
+    that is roughly 75\% smaller reduces the arithmetic intensity of the final
+    lm_head projection layer by up to 75\%, improving token-generation speed on
+    resource-constrained CPUs and edge NPUs.
+  - Codebase & Edge Independence: By outputting standard, structurally valid
+    Hugging Face files, the resulting compressed model can be loaded natively by
+    downstream optimization toolkits (such as OpenVINO or ONNX Runtime) without
+    requiring custom code layers.
+
+4. How to Use
+
+Setup & Dependencies
+
+Ensure your python environment has the required libraries installed:
+
+pip install torch safetensors transformers
+
+Directory Structure
+
+Place your original downloaded weights for dots.ocr-1.5 in a source directory,
+and designate a destination directory for the pruned output. Example:
+
+./weights/
+├── dots_ocr_1_5/                  <-- Original weights directory
+│   ├── config.json
+│   ├── tokenizer.json
+│   ├── model-00001-of-00002.safetensors
+│   └── ...
+└── dots_ocr_1_5_pruned/           <-- Target destination for pruned output
+
+Execution
+
+Run the script from your terminal:
+
+python Surgical_Vocab_Reduction_for_QWEN2B_DotOCRv1.5_V1.0untested.py
+
+Upon successful execution, the script will output the pruned weights, tokenizer
+structures, and updated configurations directly into your target directory. You
+can then load, quantize (e.g., using OpenVINO for Intel NPUs), or test the
+compressed model.
+
+5. Disclaimer & Untested Notice
+
+Please note that this script is marked as V1.0 Untested.
+
+While designed with strong safeguards—such as preserving all special layout
+coordinates, maintaining the byte-fallback system, and pruning the BPE merge
+sequences—empirical validation is required.
+
+Users are encouraged to run evaluations and verification benchmarks (such as
+measuring layout extraction accuracy, tokenization rate, and logit output
+accuracy) on their pruned models to confirm that model performance meets
+expectations for their targeted language profiles.
+
+"""
+
+import os
+import re
+import json
+import glob
+import shutil
+import torch
+from pathlib import Path
+
+# --- Configuration ---
+MODEL_DIR = r"./weights/dots_ocr_1_5"  # Source model folder
+OUTPUT_DIR = r"./weights/dots_ocr_1_5_pruned"  # Pruned output destination
+
+# Unwanted Unicode blocks to target for deletion
+UNWANTED_BLOCKS = [
+    (0x4E00, 0x9FFF),   # CJK Unified Ideographs (Chinese)
+    (0x3400, 0x4DBF),   # CJK Unified Ideographs Extension A
+    (0x0400, 0x04FF),   # Cyrillic
+    (0x0600, 0x06FF),   # Arabic
+    (0xAC00, 0xD7AF),   # Hangul Syllables (Korean)
+    (0x3040, 0x30FF),   # Hiragana & Katakana (Japanese)
+    (0x1100, 0x11FF),   # Hangul Jamo
+    (0x0590, 0x05FF),   # Hebrew
+    (0x0900, 0x097F),   # Devanagari (Hindi)
+    (0x0A00, 0x0A7F),   # Gurmukhi
+    (0x0B00, 0x0B7F),   # Oriya
+    (0x0E00, 0x0E7F),   # Thai
+]
+
+# Keys used for embeddings in Hugging Face models
+EMBED_KEYS = [
+    "model.embed_tokens.weight",       # Llama/Qwen standard input embeddings
+    "transformer.wte.weight",          # Alternate GPT/Qwen variant keys
+]
+LM_HEAD_KEYS = [
+    "lm_head.weight",                  # Output LM Head mapping
+]
+
+def should_keep_token(token_str: str) -> bool:
+    """
+    Evaluates whether a token should be retained based on character validation rules.
+    """
+    # 1. ALWAYS preserve special system & layout coordinate tokens
+    if token_str.startswith("<|") and token_str.endswith("|>"):
+        return True
+    
+    # 2. ALWAYS keep single byte/single character tokens to prevent breaking byte fallbacks
+    if len(token_str) == 1:
+        return True
+    
+    # 3. ALWAYS preserve raw byte fallbacks (e.g., <0xAF>)
+    if re.match(r"^<0x[0-9A-Fa-f]{2}>$", token_str):
+        return True
+    
+    # 4. Check characters against unwanted Unicode script ranges
+    for char in token_str:
+        codepoint = ord(char)
+        for start, end in UNWANTED_BLOCKS:
+            if start <= codepoint <= end:
+                return False
+                
+    return True
+
+def prune_tokenizer_json(src_path: Path, dest_path: Path):
+    """
+    Parses, filters, re-indexes, and saves the HF fast tokenizer configurations.
+    """
+    print(f"Reading tokenizer config: {src_path.name}")
+    with open(src_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+        
+    model_data = data.get("model", {})
+    if not model_data or model_data.get("type") != "BPE":
+        print("[Warning] Tokenizer is not BPE. Skipping structural modifications.")
+        return None, None
+        
+    vocab = model_data.get("vocab", {})
+    merges = model_data.get("merges", [])
+    
+    # Filter the vocabulary
+    kept_vocab = {}
+    kept_indices = []
+    
+    # Determine which tokens to retain
+    for token_str, token_id in vocab.items():
+        if should_keep_token(token_str):
+            kept_vocab[token_str] = token_id
+            kept_indices.append(token_id)
+            
+    # Re-index remaining tokens sequentially
+    kept_indices.sort()
+    old_to_new_id = {old_id: new_idx for new_idx, old_id in enumerate(kept_indices)}
+    
+    new_vocab = {token_str: old_to_new_id[old_id] for token_str, old_id in kept_vocab.items()}
+    
+    # Filter merges to avoid referencing deleted BPE components
+    new_merges = []
+    for merge_rule in merges:
+        parts = merge_rule.split(" ")
+        if len(parts) == 2:
+            left, right = parts[0], parts[1]
+            joined = left + right
+            # If sub-components are present in the vocab, retain the merge
+            if left in new_vocab and right in new_vocab and joined in new_vocab:
+                new_merges.append(merge_rule)
+                
+    # Update Tokenizer structural dictionaries
+    data["model"]["vocab"] = new_vocab
+    data["model"]["merges"] = new_merges
+    
+    # Re-index "added_tokens"
+    if "added_tokens" in data:
+        new_added = []
+        for token in data["added_tokens"]:
+            old_id = token.get("id")
+            if old_id in old_to_new_id:
+                token["id"] = old_to_new_id[old_id]
+                new_added.append(token)
+        data["added_tokens"] = new_added
+        
+    # Save the modified JSON configuration
+    with open(dest_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
+        
+    print(f"Pruned tokenizer vocabulary: {len(vocab)} -> {len(new_vocab)} tokens.")
+    return kept_indices, old_to_new_id
+
+def prune_legacy_vocab_files(src_dir: Path, dest_dir: Path, old_to_new_id: dict):
+    """
+    Prunes companion files like vocab.json or merges.json if present.
+    """
+    vocab_json_path = src_dir / "vocab.json"
+    if vocab_json_path.exists():
+        with open(vocab_json_path, "r", encoding="utf-8") as f:
+            vocab = json.load(f)
+        new_vocab = {}
+        for token, old_id in vocab.items():
+            if old_id in old_to_new_id:
+                new_vocab[token] = old_to_new_id[old_id]
+        with open(dest_dir / "vocab.json", "w", encoding="utf-8") as f:
+            json.dump(new_vocab, f, indent=2, ensure_ascii=False)
+            
+    merges_json_path = src_dir / "merges.json"
+    if merges_json_path.exists():
+        with open(merges_json_path, "r", encoding="utf-8") as f:
+            merges = json.load(f)
+        new_merges = []
+        for rule in merges:
+            parts = rule.split(" ")
+            if len(parts) == 2 and parts[0] in new_vocab and parts[1] in new_vocab:
+                new_merges.append(rule)
+        with open(dest_dir / "merges.json", "w", encoding="utf-8") as f:
+            json.dump(new_merges, f, indent=2, ensure_ascii=False)
+
+def update_model_configs(src_dir: Path, dest_dir: Path, new_vocab_size: int, old_to_new_id: dict):
+    """
+    Updates configuration variables and targets sequence tokens inside config and generation metadata files.
+    """
+    config_path = src_dir / "config.json"
+    if config_path.exists():
+        with open(config_path, "r", encoding="utf-8") as f:
+            config = json.load(f)
+        config["vocab_size"] = new_vocab_size
+        with open(dest_dir / "config.json", "w", encoding="utf-8") as f:
+            json.dump(config, f, indent=2)
+            
+    # Remap base configuration targets inside Generation configuration lists
+    gen_config_path = src_dir / "generation_config.json"
+    if gen_config_path.exists():
+        with open(gen_config_path, "r", encoding="utf-8") as f:
+            gconfig = json.load(f)
+        for key in ["bos_token_id", "eos_token_id", "pad_token_id"]:
+            if key in gconfig:
+                val = gconfig[key]
+                if isinstance(val, int) and val in old_to_new_id:
+                    gconfig[key] = old_to_new_id[val]
+                elif isinstance(val, list):
+                    gconfig[key] = [old_to_new_id[v] for v in val if v in old_to_new_id]
+        with open(dest_dir / "generation_config.json", "w", encoding="utf-8") as f:
+            json.dump(gconfig, f, indent=2)
+
+def prune_safetensors_weights(src_dir: Path, dest_dir: Path, kept_indices: list):
+    """
+    Loads Safetensors parameters dynamically, prunes target spatial dimensions, 
+    and writes optimized bin components.
+    """
+    from safetensors import safe_open
+    from safetensors.torch import save_file
+    
+    indices_tensor = torch.tensor(kept_indices, dtype=torch.long)
+    safetensor_files = glob.glob(str(src_dir / "*.safetensors"))
+    
+    for file_path in safetensor_files:
+        p_path = Path(file_path)
+        print(f"Processing weights file: {p_path.name}")
+        
+        tensors = {}
+        modified = False
+        
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            for k in f.keys():
+                tensor_val = f.get_tensor(k)
+                
+                # Check for input or output embedding weights
+                if k in EMBED_KEYS or k in LM_HEAD_KEYS:
+                    print(f"  -> Pruning target embedding layer: '{k}'")
+                    print(f"     Original dimensions: {list(tensor_val.shape)}")
+                    tensor_val = tensor_val[indices_tensor]
+                    print(f"     New dimensions:      {list(tensor_val.shape)}")
+                    modified = True
+                    
+                tensors[k] = tensor_val
+                
+        # Save output
+        out_file = dest_dir / p_path.name
+        save_file(tensors, str(out_file))
+        
+    return modified
+
+def prune_pytorch_bin_weights(src_dir: Path, dest_dir: Path, kept_indices: list):
+    """
+    Alternative handler for models using legacy PyTorch binary format weights files.
+    """
+    indices_tensor = torch.tensor(kept_indices, dtype=torch.long)
+    bin_files = glob.glob(str(src_dir / "*.bin"))
+    
+    for file_path in bin_files:
+        p_path = Path(file_path)
+        print(f"Processing weights file: {p_path.name}")
+        
+        state_dict = torch.load(file_path, map_location="cpu")
+        modified = False
+        
+        for k in list(state_dict.keys()):
+            if k in EMBED_KEYS or k in LM_HEAD_KEYS:
+                print(f"  -> Pruning target embedding layer: '{k}'")
+                print(f"     Original dimensions: {list(state_dict[k].shape)}")
+                state_dict[k] = state_dict[k][indices_tensor]
+                print(f"     New dimensions:      {list(state_dict[k].shape)}")
+                modified = True
+                
+        torch.save(state_dict, dest_dir / p_path.name)
+        
+    return modified
+
+def copy_miscellaneous_files(src_dir: Path, dest_dir: Path):
+    """
+    Copies companion weights files, structural classes, and configuration wrappers.
+    """
+    extensions = ["*.py", "*.png", "*.txt", "*.md", "preprocessor_config.json", "processor_config.json"]
+    for ext in extensions:
+        for f in glob.glob(str(src_dir / ext)):
+            shutil.copy(f, dest_dir)
+
+def main():
+    print("=" * 80)
+    print("  Surgical Vocab and Weight Parameter Reducer")
+    print("=" * 80)
+    
+    src_dir = Path(MODEL_DIR)
+    dest_dir = Path(OUTPUT_DIR)
+    
+    if not src_dir.exists():
+        print(f"[Error] Source model folder path does not exist: {src_dir}")
+        return
+        
+    dest_dir.mkdir(parents=True, exist_ok=True)
+    
+    # 1. Prune the tokenizer vocabulary mapping
+    tokenizer_json = src_dir / "tokenizer.json"
+    if not tokenizer_json.exists():
+        print(f"[Error] 'tokenizer.json' was not found in: {src_dir}")
+        return
+        
+    kept_indices, old_to_new_id = prune_tokenizer_json(tokenizer_json, dest_dir / "tokenizer.json")
+    if kept_indices is None:
+        print("[Error] Failed to prune tokenizer.")
+        return
+        
+    new_vocab_size = len(kept_indices)
+    
+    # 2. Prune alternative/legacy files if present
+    prune_legacy_vocab_files(src_dir, dest_dir, old_to_new_id)
+    
+    # 3. Process structural weights
+    print("\nStarting weight parameter pruning...")
+    has_safetensors = len(glob.glob(str(src_dir / "*.safetensors"))) > 0
+    
+    if has_safetensors:
+        prune_safetensors_weights(src_dir, dest_dir, kept_indices)
+    else:
+        prune_pytorch_bin_weights(src_dir, dest_dir, kept_indices)
+        
+    # 4. Save metadata settings adjustments
+    update_model_configs(src_dir, dest_dir, new_vocab_size, old_to_new_id)
+    
+    # 5. Bring over vision projections, license texts, and system processors
+    copy_miscellaneous_files(src_dir, dest_dir)
+    
+    print("\n" + "=" * 80)
+    print("Reduction pipeline completed successfully.")
+    print(f"  - Original Vocab Size (approx): {151936 if '1.5' in MODEL_DIR else 'Unknown'}")
+    print(f"  - New Pruned Vocab Size:        {new_vocab_size}")
+    print(f"  - Reduced Weight Files Saved:   {dest_dir}")
+    print("=" * 80)
+
+if __name__ == "__main__":
+    main()