Spaces:
Sleeping
Sleeping
File size: 3,533 Bytes
d1ecac0 | 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 | """The project's CNN + Transformer as a :class:`Captioner`."""
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol
from captioning.baselines.base import Captioner, CaptionerIdentity
from captioning.config.schema import AppConfig, ComparedModelConfig
if TYPE_CHECKING:
from captioning.inference.predictor import DecodeStrategy
class _ModelSettings(Protocol):
max_length: int
class _PredictorConfig(Protocol):
@property
def model(self) -> _ModelSettings: ...
class PathPredictor(Protocol):
"""The part of ``CaptionPredictor`` the adapter relies on."""
decode_strategy: str
beam_width: int
length_penalty: float
repetition_penalty: float
no_repeat_ngram_size: int
@property
def config(self) -> _PredictorConfig: ...
def predict_path(self, image_path: str | Path) -> str: ...
class CNNCaptioner(Captioner):
"""Wraps a loaded ``CaptionPredictor`` without changing how it decodes.
Each image goes through ``predictor.predict_path``, one at a time, exactly
as ``scripts/evaluate.py`` does. The recorded decode settings follow
``run_meta.json``: beam-only parameters are ``None`` for greedy decoding.
"""
def __init__(self, predictor: PathPredictor, model: ComparedModelConfig) -> None:
beam = predictor.decode_strategy == "beam"
super().__init__(
CaptionerIdentity(
model_id=model.model_id,
hub_repo=model.hub_repo,
revision=model.revision,
decode_settings={
"decode_strategy": predictor.decode_strategy,
"beam_width": predictor.beam_width if beam else None,
"length_penalty": predictor.length_penalty if beam else None,
"repetition_penalty": predictor.repetition_penalty,
"no_repeat_ngram_size": predictor.no_repeat_ngram_size,
"max_length": predictor.config.model.max_length,
},
)
)
self._predictor = predictor
@classmethod
def from_artifacts(
cls,
weights_path: str | Path,
tokenizer_dir: str | Path,
config: AppConfig,
*,
decode_strategy: DecodeStrategy | None = None,
beam_width: int | None = None,
length_penalty: float | None = None,
repetition_penalty: float | None = None,
no_repeat_ngram_size: int | None = None,
) -> CNNCaptioner:
"""Load the checkpoint with ``CaptionPredictor.from_artifacts``.
Decoding arguments are passed through unchanged, so unset ones fall
back to ``config.serve`` just as they do for ``scripts/evaluate.py``.
The identity comes from ``config.compare.cnn``. TensorFlow is imported
here, not when this module is imported.
"""
from captioning.inference.predictor import CaptionPredictor
predictor: Any = CaptionPredictor.from_artifacts(
weights_path,
tokenizer_dir,
config,
decode_strategy=decode_strategy,
beam_width=beam_width,
length_penalty=length_penalty,
repetition_penalty=repetition_penalty,
no_repeat_ngram_size=no_repeat_ngram_size,
)
return cls(predictor, config.compare.cnn)
def _raw_captions(self, image_paths: list[Path]) -> list[str]:
return [self._predictor.predict_path(path) for path in image_paths]
|