File size: 6,241 Bytes
1727faf b66671e 1727faf | 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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """BaguettotronVLM processor — self-contained for HuggingFace Hub."""
from __future__ import annotations
import json
import os
from pathlib import Path
from PIL import Image
from transformers import CLIPImageProcessor, ProcessorMixin
from transformers import PreTrainedTokenizerFast
NUM_VISUAL_TOKENS = 256
IMAGE_TOKEN = "<image>"
IMAGE_END_TOKEN = "</image>"
def _assistant_prefix(chat_style: str, enable_thinking: bool | None) -> str:
"""Return the assistant-turn content prefix for a generation prompt.
chat_style is the default baked into the repo at publish time:
- "base": stage 1 — no think tokens
- "answer": stage 2 — pre-fill </think> so the model skips reasoning
- "think": stage 3 — pre-fill <think> to trigger reasoning traces
enable_thinking overrides chat_style at call-time (stage 3 models can
toggle thinking on/off dynamically):
- None → keep chat_style default
- True → "<think>\n"
- False → "</think>\n"
"""
if enable_thinking is True:
return "<think>\n"
if enable_thinking is False:
return "</think>\n"
if chat_style == "think":
return "<think>\n"
if chat_style == "answer":
return "</think>\n"
return ""
class BaguettotronVLMProcessor(ProcessorMixin):
"""
Wraps CLIPImageProcessor + Baguettotron tokenizer.
Expands a single <image> placeholder into NUM_VISUAL_TOKENS consecutive
<image> token IDs so the model's forward() can replace them with ViT
features. Builds Qwen-style chat prompts with stage-appropriate
assistant prefixes.
Load with::
from transformers import AutoProcessor
processor = AutoProcessor.from_pretrained(
"andreagemelli/baguettotron-vision-vqa",
trust_remote_code=True,
)
"""
attributes = ["image_processor", "tokenizer"]
image_processor_class = "CLIPImageProcessor"
tokenizer_class = "AutoTokenizer"
def __init__(
self,
image_processor: CLIPImageProcessor,
tokenizer: PreTrainedTokenizerFast,
num_visual_tokens: int = NUM_VISUAL_TOKENS,
chat_style: str = "answer",
):
super().__init__(image_processor, tokenizer)
self.num_visual_tokens = num_visual_tokens
self.chat_style = chat_style
raw_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
self.image_token_id: int = raw_id if isinstance(raw_id, int) else int(raw_id[0])
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): # type: ignore[override]
processor = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
# Re-register special tokens (some tokenizers drop them on save/load)
processor.tokenizer.add_special_tokens(
{"additional_special_tokens": [IMAGE_TOKEN, IMAGE_END_TOKEN]}
)
raw_id = processor.tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
processor.image_token_id = raw_id if isinstance(raw_id, int) else int(raw_id[0])
# Load chat_style from config.json (written by push_to_hub per stage)
chat_style = "answer"
try:
if os.path.isdir(pretrained_model_name_or_path):
cfg_path = Path(pretrained_model_name_or_path) / "config.json"
if cfg_path.exists():
chat_style = json.loads(cfg_path.read_text()).get(
"chat_style", chat_style
)
else:
from huggingface_hub import hf_hub_download
cfg_path = hf_hub_download(
repo_id=pretrained_model_name_or_path, filename="config.json"
)
chat_style = json.loads(Path(cfg_path).read_text()).get(
"chat_style", chat_style
)
except Exception:
pass
processor.chat_style = chat_style
return processor
def _format_messages(
self,
messages: list[dict],
add_generation_prompt: bool,
enable_thinking: bool | None,
) -> str:
parts = [
f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>" for m in messages
]
text = "\n".join(parts)
if add_generation_prompt:
prefix = _assistant_prefix(self.chat_style, enable_thinking)
text = f"{text}\n<|im_start|>assistant\n{prefix}"
return text
def __call__(
self,
text: str | None = None,
messages: list[dict] | None = None,
image: Image.Image | None = None,
return_tensors: str = "pt",
add_generation_prompt: bool = True,
enable_thinking: bool | None = None,
) -> dict:
"""Tokenise text (or a messages list) and optionally preprocess an image.
Args:
text: raw prompt string (with a single <image> placeholder).
messages: alternative to text — list of chat dicts with "role"/"content".
image: PIL image to preprocess (optional).
add_generation_prompt: append an <|im_start|>assistant\n prefix.
enable_thinking: override chat_style for this call.
- None: keep the stage default (chat_style)
- True: pre-fill <think>\n (stage 3 reasoning mode)
- False: pre-fill </think>\n (stage 2 / stage 3 no-think mode)
"""
if messages is not None:
text = self._format_messages(
messages, add_generation_prompt, enable_thinking
)
if text is None:
raise ValueError("Provide either text or messages.")
expanded = text.replace(IMAGE_TOKEN, IMAGE_TOKEN * self.num_visual_tokens, 1)
enc = self.tokenizer(
expanded, return_tensors=return_tensors, add_special_tokens=False
)
result = {
"input_ids": enc["input_ids"],
"attention_mask": enc["attention_mask"],
}
if image is not None:
pv = self.image_processor(images=image, return_tensors=return_tensors)
result["pixel_values"] = pv.pixel_values
else:
result["pixel_values"] = None
return result
|