| """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): |
| processor = super().from_pretrained(pretrained_model_name_or_path, **kwargs) |
| |
| 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]) |
|
|
| |
| 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 |
|
|
