File size: 2,309 Bytes
90316bb | 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 | """Inference for exported BERT-CNN and ModernBERT-CNN relatedness models."""
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 RelatednessScorer:
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"]
@torch.inference_mode()
def score(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")
if batch_size < 1:
raise ValueError("batch_size must be positive")
scores = []
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)
scores.extend(self.model.probabilities(self.model(**batch)).float().cpu().tolist())
if not all(0 <= score <= 1 for score in scores):
raise RuntimeError("Model produced invalid scores")
return scores
|