| """Pretrained Hugging Face captioning baselines as a :class:`Captioner`. |
| |
| ``torch`` and ``transformers`` (the optional ``[hf]`` extra) are imported only |
| inside :meth:`HFCaptioner.load`, through ``importlib``, so importing this |
| module never loads them (ADR-019). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import importlib |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| from PIL import Image |
|
|
| from captioning.baselines.base import Captioner, CaptionerIdentity |
| from captioning.config.schema import BaselineDecodeConfig, ComparedModelConfig |
|
|
| HF_INSTALL_HINT = 'pip install -e ".[hf]"' |
|
|
|
|
| class MissingHFDependencyError(ImportError): |
| """A Hugging Face baseline was used without the ``[hf]`` extra installed.""" |
|
|
|
|
| @dataclass(frozen=True) |
| class _Loaded: |
| torch: Any |
| dtype: Any |
| image_processor: Any |
| tokenizer: Any |
| model: Any |
|
|
|
|
| class HFCaptioner(Captioner): |
| """One Hugging Face image-to-text model, pinned to an exact revision. |
| |
| BLIP-base, ViT-GPT2 and GIT-base-coco all load through |
| ``AutoImageProcessor``, ``AutoTokenizer`` and ``AutoModelForVision2Seq``. |
| Generation uses only the protocol's decode settings, passed explicitly to |
| ``generate()``, with no text prompt (``EVAL_METHODOLOGY.md`` § 8.4). |
| |
| Construction is cheap and needs no ``[hf]`` packages. Weights load on |
| :meth:`load`, or on the first :meth:`caption` call. |
| """ |
|
|
| def __init__( |
| self, |
| model: ComparedModelConfig, |
| decode: BaselineDecodeConfig, |
| *, |
| device: str = "cpu", |
| ) -> None: |
| super().__init__( |
| CaptionerIdentity( |
| model_id=model.model_id, |
| hub_repo=model.hub_repo, |
| revision=model.revision, |
| decode_settings=decode.model_dump(), |
| ) |
| ) |
| self._model = model |
| self._decode = decode |
| self._device = device |
| self._loaded: _Loaded | None = None |
|
|
| @property |
| def generate_kwargs(self) -> dict[str, Any]: |
| """The exact keyword arguments passed to ``model.generate()``.""" |
| return { |
| "num_beams": self._decode.num_beams, |
| "do_sample": self._decode.do_sample, |
| "max_new_tokens": self._decode.max_new_tokens, |
| "repetition_penalty": self._decode.repetition_penalty, |
| "no_repeat_ngram_size": self._decode.no_repeat_ngram_size, |
| } |
|
|
| def load(self) -> None: |
| """Load the processor, tokenizer and model at the pinned revision. |
| |
| Raises: |
| MissingHFDependencyError: If ``torch`` or ``transformers`` is not |
| installed. |
| """ |
| if self._loaded is not None: |
| return |
| torch, transformers = _import_hf() |
| repo, revision = self._model.hub_repo, self._model.revision |
| dtype = getattr(torch, self._decode.precision) |
| image_processor = transformers.AutoImageProcessor.from_pretrained(repo, revision=revision) |
| tokenizer = transformers.AutoTokenizer.from_pretrained(repo, revision=revision) |
| model = transformers.AutoModelForVision2Seq.from_pretrained( |
| repo, revision=revision, torch_dtype=dtype |
| ) |
| model = model.to(self._device) |
| model.eval() |
| self._loaded = _Loaded(torch, dtype, image_processor, tokenizer, model) |
|
|
| def _raw_captions(self, image_paths: list[Path]) -> list[str]: |
| self.load() |
| loaded = self._loaded |
| assert loaded is not None |
| images = [_open_rgb(path) for path in image_paths] |
| pixel_values = loaded.image_processor(images=images, return_tensors="pt").pixel_values |
| pixel_values = pixel_values.to(self._device, dtype=loaded.dtype) |
| with loaded.torch.inference_mode(): |
| output_ids = loaded.model.generate(pixel_values=pixel_values, **self.generate_kwargs) |
| return list(loaded.tokenizer.batch_decode(output_ids, skip_special_tokens=True)) |
|
|
|
|
| def _import_hf() -> tuple[Any, Any]: |
| try: |
| torch = importlib.import_module("torch") |
| transformers = importlib.import_module("transformers") |
| except ImportError as exc: |
| raise MissingHFDependencyError( |
| "Hugging Face baselines need the optional [hf] extra (transformers, torch). " |
| f"Install it with: {HF_INSTALL_HINT}" |
| ) from exc |
| return torch, transformers |
|
|
|
|
| def _open_rgb(path: Path) -> Image.Image: |
| with Image.open(path) as image: |
| return image.convert("RGB") |
|
|
