Publish ModernBERT-CNN with concise usage, training details and own evaluation metrics
90316bb verified | """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"] | |
| 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 | |