Update Wall-X to 1.1.0 (#104)

2026-06-15 11:40:00 +08:00
parent e23a586846
commit 72834e7de5
200 changed files with 33916 additions and 16771 deletions
@@ -1,26 +1,107 @@
-from transformers import AutoProcessor
-import os
+#!/usr/bin/env python3
+"""Merge Wall-X action tokens into a Qwen2.5-VL processor tokenizer."""

-processor_path = "/path/to/Qwen2.5-VL-3B-Instruct"
-action_tokenizer_path = "/path/to/fast"
-use_fast_tokenizer = True
+from __future__ import annotations

-processor = AutoProcessor.from_pretrained(processor_path, use_fast=True)
-processor.tokenizer.padding_side = "left"
+import argparse
+from pathlib import Path

-action_tokenizer = AutoProcessor.from_pretrained(
-    action_tokenizer_path, trust_remote_code=True
-)

-new_tokens = ["<|propri|>", "<|action|>"]
-new_tokens += [f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)]
-num_added_tokens = processor.tokenizer.add_tokens(new_tokens)
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Create a Wall-X processor directory by adding FAST action tokens "
+            "to a Qwen2.5-VL processor tokenizer."
+        )
+    )
+    parser.add_argument(
+        "--processor-path",
+        required=True,
+        help="Base Qwen2.5-VL processor directory or Hugging Face repo id.",
+    )
+    parser.add_argument(
+        "--action-tokenizer-path",
+        required=True,
+        help="FAST/action tokenizer processor directory or Hugging Face repo id.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        required=True,
+        help="Directory where the merged processor will be written.",
+    )
+    parser.add_argument(
+        "--use-fast",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Use the fast tokenizer implementation when loading the base processor.",
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Allow custom code when loading the action tokenizer processor.",
+    )
+    return parser.parse_args()

-begin_idx_token = "<|action_token_0|>"
-token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token)
-processor.tokenizer.init_kwargs["action_token_start_index"] = token_id
-processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_tokenizer.vocab_size

-new_tokenizer_dir = "/path/to/new_tokenizer"
-os.makedirs(new_tokenizer_dir, exist_ok=True)
-processor.save_pretrained(new_tokenizer_dir)
+def _resolve_action_vocab_size(action_processor) -> int:
+    vocab_size = getattr(action_processor, "vocab_size", None)
+    if vocab_size is None and hasattr(action_processor, "tokenizer"):
+        vocab_size = getattr(action_processor.tokenizer, "vocab_size", None)
+    if vocab_size is None:
+        raise AttributeError(
+            "Could not determine action tokenizer vocab size from the loaded processor."
+        )
+    return int(vocab_size)
+
+
+def merge_tokenizer(
+    *,
+    processor_path: str,
+    action_tokenizer_path: str,
+    output_dir: str,
+    use_fast: bool,
+    trust_remote_code: bool,
+) -> None:
+    from transformers import AutoProcessor
+
+    processor = AutoProcessor.from_pretrained(processor_path, use_fast=use_fast)
+    processor.tokenizer.padding_side = "left"
+
+    action_processor = AutoProcessor.from_pretrained(
+        action_tokenizer_path,
+        trust_remote_code=trust_remote_code,
+    )
+    action_vocab_size = _resolve_action_vocab_size(action_processor)
+
+    new_tokens = ["<|propri|>", "<|action|>"]
+    new_tokens += [f"<|action_token_{i}|>" for i in range(action_vocab_size)]
+    num_added_tokens = processor.tokenizer.add_tokens(new_tokens)
+
+    begin_idx_token = "<|action_token_0|>"
+    token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token)
+    processor.tokenizer.init_kwargs["action_token_start_index"] = token_id
+    processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_vocab_size
+
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+    processor.save_pretrained(output_path)
+
+    print(f"Saved merged processor to {output_path}")
+    print(f"Added {num_added_tokens} tokenizer tokens")
+    print(f"action_token_start_index={token_id}")
+    print(f"action_token_vocab_size={action_vocab_size}")
+
+
+def main() -> None:
+    args = parse_args()
+    merge_tokenizer(
+        processor_path=args.processor_path,
+        action_tokenizer_path=args.action_tokenizer_path,
+        output_dir=args.output_dir,
+        use_fast=args.use_fast,
+        trust_remote_code=args.trust_remote_code,
+    )
+
+
+if __name__ == "__main__":
+    main()