"""MetaEncoder v1: encode inputs, retrieve from a corpus, and decide Jev-style questions with one embedding model. from modeling_metaencoder import MetaEncoder model = MetaEncoder.from_pretrained("") # 1. encode: one L2-normalised embedding per item emb = model.encode(["a red bicycle", {"image": img}], instruction="Represent the given image.") # 2. retrieve: rank a corpus (items, or embeddings from `encode`) against a query hits = model.retrieve("Which photo shows a bicycle?", corpus, top_k=5) # 3. decide: Jev-style questions about a state model.decide({"query": "used iphone 15"}, {"shopping": {"type": "noul", "instructions": "Is the user shopping for a product?"}}) # -> {"questions": {"shopping": {"value": True, "probabilities": {"false": ..., "true": ...}, ...}}} Built only on the public `transformers` Muse Glimmer classes (`AutoModel` -> `MuseGlimmerModel`, `AutoProcessor` -> `MuseGlimmerProcessor`). The embedding is the final-norm hidden state of the input's closing `<|eot|>` token, projected by `embedding_head.safetensors` and L2-normalised. """ from __future__ import annotations import json import math import os import warnings from collections import defaultdict from dataclasses import dataclass from typing import Any, Mapping, Sequence, Union import torch import torch.nn.functional as F from huggingface_hub import hf_hub_download from PIL import Image from safetensors.torch import load_file from transformers import AutoModel, AutoProcessor MAX_TRAINED_LENGTH = 16384 DEFAULT_BATCH_SIZE = 8 MAX_IMAGE_TOKENS = 4096 VIDEO_FRAMES = 8 TASK_INSTRUCTION = "Read the task definition and select the option it calls for." SCENARIO_INSTRUCTION = "Answer the question about the scenario by selecting the correct option." TRAINED_TEMPERATURE = 0.02 # Options of at most this many words are labels whose meaning comes from the task, so their candidates carry the # task as their instruction (training kept the task only up to MAX_TASK_CHARS). LABEL_MAX_WORDS = 5 MAX_TASK_CHARS = 600 MAX_OPTIONS = 26 # Trained option order of yes/no ("noul") questions and the Jev answer key of each. NOUL_OPTIONS = (("No", "false"), ("Yes", "true")) _MAX_CACHED_CANDIDATES = 4096 Item = Union[str, Image.Image, Mapping[str, Any]] _ITEM_KEYS = {"text", "instruction", "image", "images", "video"} @dataclass class Hit: """One retrieved corpus entry.""" index: int # position in the corpus passed in score: float # cosine similarity to the query, in [-1, 1] item: Any # the corpus item itself, or None if embeddings were passed def __repr__(self) -> str: return f"Hit(index={self.index}, score={self.score:.4f})" @dataclass(frozen=True) class Question: """A parsed Jev question: what the encoder sees (options) and what the answer reports (keys).""" name: str type: str # "choice", "noul" or "score" instructions: str options: tuple[str, ...] # option texts in prompt order keys: tuple[str, ...] # Jev answer key per option: criteria key, "false"/"true", or "0".."n-1" descriptions: tuple[str, ...] # per option, "" if none def _sample_frames(frames: Sequence[Any]) -> list[Any]: if len(frames) <= VIDEO_FRAMES: return list(frames) return [frames[round(i * (len(frames) - 1) / (VIDEO_FRAMES - 1))] for i in range(VIDEO_FRAMES)] def _as_items(inputs: Item | Sequence[Item], instruction: str | None) -> list[dict]: """Normalise any accepted input shape into one dict per embedding.""" if isinstance(inputs, (str, Image.Image, Mapping)): inputs = [inputs] items = [] for i, raw in enumerate(inputs): if isinstance(raw, str): raw = {"text": raw} elif isinstance(raw, Image.Image): raw = {"image": raw} elif not isinstance(raw, Mapping): raise TypeError(f"item {i} is {type(raw).__name__}; expected str, PIL.Image or a dict") unknown = set(raw) - _ITEM_KEYS if unknown: raise ValueError(f"item {i} has unknown key(s) {sorted(unknown)}; allowed: {sorted(_ITEM_KEYS)}") images = [] for key in ("image", "images", "video"): value = raw.get(key) if value is None: continue value = list(value) if isinstance(value, (list, tuple)) else [value] images.extend(_sample_frames(value) if key == "video" else value) items.append({ "text": raw.get("text") or "", "instruction": raw.get("instruction", instruction) or "", "images": [im.convert("RGB") for im in images if im is not None], }) return items def parse_question(name: str, spec: Mapping[str, Any]) -> Question: """Validate one Jev question `{"type", "instructions", "criteria"}`. choice: `criteria` maps each option key to its description (may be empty). noul: optional `{"true": ..., "false": ...}` descriptions. score: `criteria` is the ordered list of level labels. """ if not isinstance(spec, Mapping) or set(spec) - {"type", "instructions", "criteria"}: raise ValueError(f"{name}: a question is {{'type', 'instructions', 'criteria'}}") kind, criteria = spec.get("type"), spec.get("criteria") instructions = " ".join(str(spec.get("instructions") or "").split()) if not instructions: raise ValueError(f"{name}: instructions must be non-empty") if kind == "noul": criteria = criteria or {} if not isinstance(criteria, Mapping) or set(criteria) - {"true", "false"}: raise ValueError(f"{name}: noul criteria may only describe 'true' and 'false'") options, keys = zip(*NOUL_OPTIONS) descriptions = [criteria.get(k) for k in keys] elif kind == "choice": if not isinstance(criteria, Mapping): raise ValueError(f"{name}: choice criteria must map option keys to descriptions") options = keys = tuple(str(k).strip() for k in criteria) descriptions = list(criteria.values()) elif kind == "score": if isinstance(criteria, (str, Mapping)) or not isinstance(criteria, Sequence): raise ValueError(f"{name}: score criteria must be the ordered list of level labels") options = tuple(str(c).strip() for c in criteria) keys = tuple(str(i) for i in range(len(options))) descriptions = [""] * len(options) else: raise ValueError(f"{name}: type must be choice, noul or score, got {kind!r}") if not 2 <= len(options) <= MAX_OPTIONS or len(set(options)) != len(options) or not all(options): raise ValueError(f"{name}: needs 2 to {MAX_OPTIONS} distinct non-empty options, got {list(options)}") descriptions = tuple("" if d is None or str(d).strip() == o else str(d).strip() for o, d in zip(options, descriptions)) return Question(name, kind, instructions, tuple(options), tuple(keys), descriptions) def render_state(state: Any) -> tuple[str, str]: """(state text, query instruction): text as is under the typed-decision instruction, JSON as a scenario.""" if state is None or isinstance(state, str): return (state or "").strip(), TASK_INSTRUCTION return json.dumps(state, ensure_ascii=False), SCENARIO_INSTRUCTION def render_options(options: Sequence[str]) -> str: """The option list of a decision query: comma-joined labels, or one "- option" line each when sentence-like.""" if any(len(o) > 40 or any(c in o for c in ",.;\n") for o in options): return "Options:\n" + "\n".join(f"- {o}" for o in options) return "Options: " + ", ".join(options) def decision_prompts(state_text: str, question: str, options: Sequence[str], descriptions: Sequence[str], instruction: str) -> tuple[dict, list[dict]]: """The query and one candidate per option of a decision, laid out as the training rows (images excluded).""" head = f"{state_text}\n\n" if state_text else "" query = {"text": f"{head}Task: {question.strip()}\n{render_options(options)}", "instruction": instruction} task = " ".join(question.split()) label_style = all(len(o.split()) <= LABEL_MAX_WORDS for o in options) cand_instruction = f"Option for the task: {task}" if label_style and len(task) <= MAX_TASK_CHARS else "" candidates = [{"text": f"{o}\n{d}" if d else o, "instruction": cand_instruction, "images": []} for o, d in zip(options, descriptions)] return query, candidates def answer(question: Question, scores: torch.Tensor, temperature: float) -> dict[str, Any]: """The Jev answer of one question from its option cosine scores.""" probs = torch.softmax(scores.double() / temperature, dim=-1) best = int(probs.argmax()) entropy = -float((probs * probs.clamp_min(1e-300).log()).sum()) if question.type == "score": value: Any = float(sum(i * p for i, p in enumerate(probs.tolist()))) elif question.type == "noul": value = question.keys[best] == "true" else: value = question.keys[best] return {"value": value, "probabilities": dict(zip(question.keys, probs.tolist())), "confidence": min(1.0, max(0.0, 1.0 - entropy / math.log(len(question.keys)))), "scores": dict(zip(question.keys, scores.tolist()))} def fit_temperature(rows: Sequence[tuple[Sequence[float], int]]) -> float: """The softmax temperature that minimises the NLL of (option scores, correct index) rows.""" grid = torch.logspace(math.log10(1e-3), 0.0, 301, dtype=torch.float64) groups: dict[int, list[tuple[Sequence[float], int]]] = defaultdict(list) for scores, label in rows: groups[len(scores)].append((scores, label)) nll = torch.zeros_like(grid) for group in groups.values(): scores = torch.tensor([s for s, _ in group], dtype=torch.float64) labels = torch.tensor([y for _, y in group]) nll += torch.stack([F.cross_entropy(scores / t, labels, reduction="sum") for t in grid]) return float(grid[int(nll.argmin())]) class MetaEncoder: """Embedding model over text, images and video (as a list of frames). Usage: model = MetaEncoder.from_pretrained("") corpus = model.encode(docs, instruction="Represent the document for retrieval.") model.retrieve(query, corpus, top_k=5) model.decide({"post": text}, {"sports": {"type": "noul", "instructions": "Is the post about sports?"}}) """ def __init__(self, model: torch.nn.Module, processor: Any, head: torch.Tensor): self.model = model self.processor = processor self.head = head # Softmax temperature of `decide`; `calibrate` refits it on labelled answers. self.temperature = TRAINED_TEMPERATURE # Option candidates depend only on the question, so repeated questions reuse their embeddings. self._candidate_cache: dict[tuple[str, str], torch.Tensor] = {} # Last-token pooling reads the final column, which is the closing <|eot|> only under left padding. self.processor.tokenizer.padding_side = "left" @classmethod def from_pretrained(cls, path: str, device_map: str | dict = "auto", dtype: torch.dtype = torch.bfloat16, attn_implementation: str = "sdpa") -> "MetaEncoder": model = AutoModel.from_pretrained(path, dtype=dtype, device_map=device_map, attn_implementation=attn_implementation) model.eval() head_file = os.path.join(path, "embedding_head.safetensors") if not os.path.exists(head_file): head_file = hf_hub_download(path, "embedding_head.safetensors") head = load_file(head_file)["embedding_head.weight"].to(model.device, dtype) return cls(model, AutoProcessor.from_pretrained(path), head) @property def embedding_dim(self) -> int: return self.head.shape[0] def encode(self, inputs: Item | Sequence[Item], instruction: str | None = None, batch_size: int = DEFAULT_BATCH_SIZE) -> torch.Tensor: """Embed items. Returns L2-normalised float32 `(n_items, embedding_dim)`. An item is a string, a PIL image, or a dict `{"text", "image", "images", "video", "instruction"}` (all optional). `instruction` describes what the embedding is for (e.g. "Find a caption for the news in the given photo.") and applies to items that carry none of their own. """ return self._encode_items(_as_items(inputs, instruction), batch_size) def retrieve(self, query: Item, corpus: Sequence[Item] | torch.Tensor, top_k: int | None = None, query_instruction: str | None = None, corpus_instruction: str | None = None, batch_size: int = DEFAULT_BATCH_SIZE) -> list[Hit]: """Rank `corpus` against `query`, best first. `corpus` is a list of items, or their embeddings from `encode` (then `corpus_instruction` is unused). Encode a large corpus once and pass the tensor to search it repeatedly. """ q = self.encode(query, instruction=query_instruction, batch_size=1) is_tensor = isinstance(corpus, torch.Tensor) c = corpus if is_tensor else self.encode(corpus, instruction=corpus_instruction, batch_size=batch_size) scores = (q @ c.to(q.device, q.dtype).T)[0] order = torch.argsort(scores, descending=True)[:top_k] originals = None if is_tensor else list(corpus) return [Hit(int(i), float(scores[i]), None if originals is None else originals[int(i)]) for i in order] def decide(self, state: Any, questions: Mapping[str, Mapping[str, Any]], images: Any = None, video: Sequence[Image.Image] | None = None) -> dict[str, Any]: """Answer Jev questions about `state`. Args: state: what the questions are about: a string, or any JSON value (rendered as a JSON scenario). questions: `{name: {"type": "choice" | "noul" | "score", "instructions": ..., "criteria": ...}}`. images: a PIL image or a list of them that belong to the state. video: the state's video as a list of frames. Returns: `{"questions": {name: {"value", "probabilities", "confidence", "scores"}}}`. `value` is the chosen criteria key (choice), a bool (noul) or the expected level index (score); `probabilities` maps each answer key to softmax(score / temperature); `confidence` is 1 - normalised entropy; `scores` are the raw cosine similarities. """ return self.decide_batch([state], questions, [images], [video])[0] def decide_batch(self, states: Sequence[Any], questions: Mapping[str, Mapping[str, Any]], images: Sequence[Any] | None = None, video: Sequence[Sequence[Image.Image] | None] | None = None, batch_size: int = DEFAULT_BATCH_SIZE) -> list[dict[str, Any]]: """`decide` for many states that share the same questions; option embeddings are computed once.""" parsed = [parse_question(name, spec) for name, spec in questions.items()] if not parsed: raise ValueError("questions must name at least one question") n = len(states) images, video = images or [None] * n, video or [None] * n if not len(images) == len(video) == n: raise ValueError("pass one images and one video entry per state") rendered = [render_state(s) for s in states] media = [_as_items({"images": im, "video": v}, None)[0]["images"] for im, v in zip(images, video)] queries, cands = [], [] for q in parsed: for (text, instruction), state_images in zip(rendered, media): query, cands_q = decision_prompts(text, q.instructions, q.options, q.descriptions, instruction) queries.append({**query, "images": state_images}) cands.append(cands_q) query_embs = self._encode_items(queries, batch_size).view(len(parsed), n, -1) results: list[dict[str, Any]] = [{"questions": {}} for _ in states] for q, embs, cands_q in zip(parsed, query_embs, cands): scores = embs @ self._encode_candidates(cands_q, batch_size).T for result, s in zip(results, scores): result["questions"][q.name] = answer(q, s, self.temperature) return results def calibrate(self, answers: Sequence[Mapping[str, Any]], labels: Sequence[Any]) -> float: """Fit `self.temperature` to labelled answers and return it; later `decide` calls use it. `answers` are per-question answers from `decide` (`result["questions"][name]`); `labels` are their correct answer keys (a criteria key, True/False for noul, or a level index for score). """ if len(answers) != len(labels) or not answers: raise ValueError("need one label per answer, and at least one answer") rows = [] for a, label in zip(answers, labels): key = ("true" if label else "false") if isinstance(label, bool) else str(label) rows.append((list(a["scores"].values()), list(a["scores"]).index(key))) self.temperature = fit_temperature(rows) return self.temperature @staticmethod def similarity(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Cosine similarity matrix `(n_a, n_b)` of unit-norm embeddings.""" return a @ b.T def _encode_items(self, items: list[dict], batch_size: int) -> torch.Tensor: if not items: return torch.empty(0, self.embedding_dim, device=self.model.device) return torch.cat([self._encode_batch(items[i:i + batch_size]) for i in range(0, len(items), batch_size)]) def _encode_candidates(self, items: list[dict], batch_size: int) -> torch.Tensor: """Encode text-only option candidates through the cache.""" cache = self._candidate_cache keys = [(it["instruction"], it["text"]) for it in items] missing = list(dict.fromkeys(k for k in keys if k not in cache)) if missing: if len(cache) + len(missing) > _MAX_CACHED_CANDIDATES: cache.clear() embs = self._encode_items([{"instruction": i, "text": t, "images": []} for i, t in missing], batch_size) cache.update(zip(missing, embs)) return torch.stack([cache[k] for k in keys]) @torch.no_grad() def _encode_batch(self, items: list[dict]) -> torch.Tensor: texts, images = [], [] for it in items: content = [{"type": "image"} for _ in it["images"]] if it["text"]: content.append({"type": "text", "text": it["text"]}) messages = [{"role": "system", "content": it["instruction"]}] if it["instruction"] else [] messages.append({"role": "user", "content": content}) texts.append(self.processor.apply_chat_template(messages, tokenize=False)) images.extend(it["images"]) inputs = self.processor(text=texts, images=images or None, padding=True, truncation=False, add_special_tokens=False, max_image_tokens=MAX_IMAGE_TOKENS, return_tensors="pt") inputs = {k: v.to(self.model.device) if hasattr(v, "to") else v for k, v in inputs.items()} if not bool((inputs["attention_mask"][:, -1] == 1).all()): raise RuntimeError("expected left padding") longest = int(inputs["attention_mask"].sum(-1).max()) if longest > MAX_TRAINED_LENGTH: warnings.warn(f"an input has {longest} tokens, longer than the {MAX_TRAINED_LENGTH} seen in training") hidden = self.model(**inputs, use_cache=False).last_hidden_state[:, -1] return F.normalize((hidden @ self.head.T).float(), dim=-1)