processing_vision_gptoss.py

"""Processor for vision-gpt-oss: Gemma-4 image processor + gpt-oss (harmony) tokenizer.

Builds the LLaVA-style visual block ``<|vis_start|> <|vis_pad|>*n <|vis_end|>`` (n =
soft tokens for the image) and renders the gpt-oss harmony chat template.
"""
from transformers.processing_utils import ProcessorMixin
from transformers.feature_extraction_utils import BatchFeature


class VisionGptOssProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(self, image_processor=None, tokenizer=None,
                 vis_start_token="<|vis_start|>", vis_pad_token="<|vis_pad|>",
                 vis_end_token="<|vis_end|>", **kwargs):
        self.vis_start_token = vis_start_token
        self.vis_pad_token = vis_pad_token
        self.vis_end_token = vis_end_token
        super().__init__(image_processor, tokenizer)

    def _vis_block(self, n):
        return self.vis_start_token + self.vis_pad_token * n + self.vis_end_token

    def __call__(self, images=None, text=None, messages=None, return_tensors="pt",
                 reasoning_effort="low", add_generation_prompt=True, **kwargs):
        img_out, block = {}, ""
        if images is not None:
            if not isinstance(images, (list, tuple)):
                images = [images]
            img_out = self.image_processor(images=images, return_tensors="pt")
            n = int(img_out["num_soft_tokens_per_image"][0])
            block = self._vis_block(n)

        if messages is not None:
            msgs = [dict(m) for m in messages]
            if block:
                for m in msgs:
                    if m.get("role") == "user":
                        m["content"] = block + "\n" + str(m["content"])
                        break
            full = self.tokenizer.apply_chat_template(
                msgs, add_generation_prompt=add_generation_prompt,
                tokenize=False, reasoning_effort=reasoning_effort)
        else:
            content = (block + "\n" + (text or "")) if block else (text or "")
            full = self.tokenizer.apply_chat_template(
                [{"role": "user", "content": content}],
                add_generation_prompt=add_generation_prompt,
                tokenize=False, reasoning_effort=reasoning_effort)

        enc = self.tokenizer(full, add_special_tokens=False, return_tensors=return_tensors)
        data = {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]}
        if images is not None:
            data["pixel_values"] = img_out["pixel_values"]
            data["image_position_ids"] = img_out["image_position_ids"]
        return BatchFeature(data)

    def batch_decode(self, *a, **k):
        return self.tokenizer.batch_decode(*a, **k)

    def decode(self, *a, **k):
        return self.tokenizer.decode(*a, **k)


__all__ = ["VisionGptOssProcessor"]