File size: 4,543 Bytes
2a190f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# processing_opencua.py
import torch
from typing import List, Dict, Any, Union
from PIL import Image
from transformers.processing_utils import ProcessorMixin, BatchFeature
from transformers import AutoTokenizer, AutoImageProcessor
PLACEHOLDER = "<|media_placeholder|>"
class OpenCUAProcessor(ProcessorMixin):
attributes = ["image_processor", "tokenizer", "image_token_id", "merge_size"]
def __init__(self, image_processor, tokenizer, image_token_id: int = 151664, merge_size: int = 2, **kwargs):
self.image_processor = image_processor
self.tokenizer = tokenizer
self.image_token_id = image_token_id
self.merge_size = getattr(image_processor, "merge_size", merge_size)
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
trust = kwargs.get("trust_remote_code", True)
# 优先用你仓库的 TikTokenV3；失败回退 AutoTokenizer（只用于初始化/占位）
try:
from tokenization_opencua import TikTokenV3
tok = TikTokenV3.from_pretrained(pretrained_model_name_or_path, trust_remote_code=trust)
except Exception:
tok = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, trust_remote_code=trust)
imgproc = AutoImageProcessor.from_pretrained(pretrained_model_name_or_path, trust_remote_code=trust)
return cls(imgproc, tok, **kwargs)
def apply_chat_template(self, messages: List[Dict[str, Any]], **kwargs) -> Union[str, List[int]]:
return self.tokenizer.apply_chat_template(messages, **kwargs)
# 下面这些方法给 HF 路径用；vLLM 初始化只需要能成功 new 出来即可
def __call__(self, *args, **kwargs) -> BatchFeature:
# 返回一个最小结构，避免被实际调用时崩溃
data = {"input_ids": torch.zeros(1, 1, dtype=torch.long)}
return BatchFeature(data=data)
# 提供给你自己脚本用的辅助（可选）
def prepare_vllm_inputs(self, messages, images, add_generation_prompt=True):
text = self.apply_chat_template(messages, tokenize=False, add_generation_prompt=add_generation_prompt)
proc = self.image_processor(images=images, return_tensors="pt")
grid = torch.as_tensor(proc["image_grid_thw"])
merge = getattr(self, "merge_size", 2)
for thw in grid:
num = int((thw[0] * thw[1] * thw[2]) // (merge ** 2))
text = text.replace(PLACEHOLDER, PLACEHOLDER * num, 1)
return text, images
# # processing_opencua.py
# from transformers import Qwen2_5_VLProcessor, AutoTokenizer, AutoImageProcessor
# class OpenCUAProcessor(Qwen2_5_VLProcessor):
# # 用字符串就行，但我们会在 from_pretrained 里手动加载，避免字符串反射
# tokenizer_class = "TikTokenV3"
# @classmethod
# def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
# # 确保 remote code 可用
# trust_remote_code = kwargs.get("trust_remote_code", False)
# # 1) 手动加载 tokenizer（会按模型目录里的 tokenizer_config.json -> TikTokenV3 + tokenization_opencua.py）
# tokenizer = AutoTokenizer.from_pretrained(
# pretrained_model_name_or_path,
# trust_remote_code=trust_remote_code,
# )
# # 2) 手动加载图像处理器（保持 Qwen2VLImageProcessor）
# image_processor = AutoImageProcessor.from_pretrained(
# pretrained_model_name_or_path,
# trust_remote_code=trust_remote_code,
# )
# # 3) 获取chat_template，如果tokenizer有的话
# chat_template = getattr(tokenizer, 'chat_template', None)
# # 4) 构造并返回 Qwen2.5-VL 的 Processor 实例，传递chat_template
# processor = cls(image_processor=image_processor, tokenizer=tokenizer, chat_template=chat_template)
# # 5) 添加vLLM需要的属性
# # 这些token ID需要与tokenizer_config.json中的定义一致
# processor.image_token = "<|media_placeholder|>" # 使用OpenCUA的媒体占位符
# processor.video_token = "<|media_placeholder|>" # 视频也使用相同的占位符
# # 添加token ID（从tokenizer_config.json中获取）
# vocab = tokenizer.get_vocab()
# processor.image_token_id = vocab.get("<|media_placeholder|>", 151664) # 默认151664
# processor.video_token_id = vocab.get("<|media_placeholder|>", 151664) # 视频使用相同ID
# return processor
|