| """Inference for the reward-fine-tuned BERT-CNN caption reranker.""" |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
| from modeling import VisualRelatednessModel, backbone_config_from_dict |
|
|
|
|
| class CaptionReranker: |
| """Pick the best of nine beam-search captions: score_k = bias[k] + weight * BERT-CNN logit.""" |
|
|
| def __init__(self, model_path, device="cpu"): |
| path = Path(model_path) |
| if not path.is_dir(): |
| from huggingface_hub import snapshot_download |
| path = Path(snapshot_download(str(model_path), allow_patterns=["config.json", "model.safetensors", "tokenizer/*"])) |
| config = json.loads((path / "config.json").read_text()) |
| self.device = torch.device(device) |
| self.model = VisualRelatednessModel(backbone_config_from_dict(config["bert_config"]), **config["model_options"]) |
| self.model.load_state_dict(load_file(str(path / "model.safetensors"))) |
| self.model.to(self.device).eval() |
| self.tokenizer = AutoTokenizer.from_pretrained(path / "tokenizer", local_files_only=True, use_fast=True) |
| self.tokenizer.padding_side = "right" |
| self.seq_len = config["model_options"]["seq_len"] |
| self.bias = config["rank_prior"]["bias"] |
| self.weight = config["rank_prior"]["weight"] |
|
|
| @torch.inference_mode() |
| def logits(self, visual_context, captions, batch_size=16): |
| if not isinstance(visual_context, str) or not visual_context.strip(): |
| raise ValueError("Provide image-context words, such as dog, ball, grass") |
| if not captions or any(not isinstance(c, str) or not c.strip() for c in captions): |
| raise ValueError("Provide at least one nonempty caption") |
| values = [] |
| for start in range(0, len(captions), batch_size): |
| selected = captions[start:start + batch_size] |
| batch = self.tokenizer([visual_context] * len(selected), selected, padding="max_length", |
| truncation="longest_first", max_length=self.seq_len, |
| return_tensors="pt").to(self.device) |
| values.extend(self.model(**batch).float().cpu().tolist()) |
| return values |
|
|
| def rerank(self, visual_context, beams): |
| """beams: the nine captions in beam-search order (best beam first). Returns (index, scores).""" |
| if len(beams) != len(self.bias): |
| raise ValueError(f"Expected exactly {len(self.bias)} captions in beam-search order") |
| scores = [b + self.weight * x for b, x in zip(self.bias, self.logits(visual_context, beams))] |
| best = max(range(len(scores)), key=lambda i: (scores[i], -i)) |
| return best, scores |
|
|
