Blame for modeling_metaencoder.py

"""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("<this repo>")

    # 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("<repo>")
        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)