simit 0.1.0__py3-none-any.whl

simit/backends/vllm.py

@@ -0,0 +1,204 @@
+"""Fast engine for standard VLMs: vLLM (continuous batching, prefix caching,
+CUDA graphs). Optional: ``pip install "simit[vllm]"``.
+
+Prompts are built with the model's own chat template (as in the transformers
+backend), so results match it up to numerics.
+"""
+
+from __future__ import annotations
+
+import itertools
+import queue
+import threading
+import traceback
+from concurrent.futures import Future
+from typing import Optional
+
+from PIL import Image
+
+from ..utils import strip_thinking, to_pil
+from .base import Backend, GenRequest, GenResult, ScoreRequest
+
+
+class VLLMBackend(Backend):
+    name = "vllm"
+    native_image_generation = False
+    verify_think = False
+
+    def __init__(self, model: str, *, gpu_memory_utilization: Optional[float] = None, max_model_len: int = 32768,
+                 max_images_per_prompt: int = 8, max_image_pixels: Optional[int] = 1024 * 1024,
+                 tensor_parallel_size: int = 1, think: bool = False, device: Optional[str] = None,
+                 max_num_seqs: int = 64, **llm_kwargs):
+        """``gpu_memory_utilization=None`` takes what is free on the GPU (minus
+        headroom), so an image generator loaded first keeps its share."""
+        import os
+        if device is not None and str(device) not in ("cuda", "cuda:0"):
+            raise ValueError("the vLLM engine runs on the first visible GPU; select GPUs with "
+                             "CUDA_VISIBLE_DEVICES (and tensor_parallel_size for several)")
+        if gpu_memory_utilization is None:
+            import torch
+            free, total = torch.cuda.mem_get_info(0)
+            gpu_memory_utilization = max(0.3, min(0.9, (free - 6 * 2**30) / total))
+        # Keep vLLM's engine core in this process: a spawned core would re-import
+        # the user's __main__ (scripts without an `if __name__ == "__main__"` guard).
+        os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+        # Portable defaults: no kernels that JIT-compile with the system nvcc
+        # (FlashInfer) or need a matching flash-attn build. Override via kwargs/env.
+        os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
+        os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")   # FP8 checkpoints: DeepGEMM JIT needs nvcc >= 12.9
+        llm_kwargs.setdefault("attention_backend", "TRITON_ATTN")
+        llm_kwargs.setdefault("mm_encoder_attn_backend", "TORCH_SDPA")
+        llm_kwargs.setdefault("gdn_prefill_backend", "triton")
+        from transformers import AutoProcessor
+        from vllm import LLM
+
+        self.model_id = model
+        self.processor = AutoProcessor.from_pretrained(model)
+        self.tokenizer = getattr(self.processor, "tokenizer", self.processor)
+        self.llm = LLM(model=model, dtype="bfloat16", gpu_memory_utilization=gpu_memory_utilization,
+                       max_model_len=max_model_len, limit_mm_per_prompt={"image": max_images_per_prompt},
+                       enable_prefix_caching=True, tensor_parallel_size=tensor_parallel_size,
+                       max_num_seqs=max_num_seqs, **llm_kwargs)
+        self.engine = self.llm.llm_engine
+        gc = self.llm.llm_engine.model_config.try_get_generation_config() or {}
+        eos = gc.get("eos_token_id")
+        eos = eos if isinstance(eos, list) else [eos] if eos is not None else []
+        self.eos_ids = {e for e in eos + [self.tokenizer.eos_token_id] if e is not None}
+        self.think = think
+        self.max_image_pixels = max_image_pixels
+        self._ids = itertools.count()
+        self._queue: "queue.Queue" = queue.Queue()
+        self._live: dict = {}
+        self._thread = threading.Thread(target=self._loop, daemon=True, name="simit-vllm")
+        self._thread.start()
+
+    # ------------------------------------------------------------- prompts
+    def _image(self, img):
+        img = to_pil(img)
+        if self.max_image_pixels and img.width * img.height > self.max_image_pixels:
+            s = (self.max_image_pixels / (img.width * img.height)) ** 0.5
+            img = img.resize((max(28, int(img.width * s)), max(28, int(img.height * s))), Image.BICUBIC)
+        return img
+
+    def _prompt(self, parts, think: bool):
+        content, images = [], []
+        for p in parts:
+            if isinstance(p, Image.Image):
+                content.append({"type": "image"})
+                images.append(self._image(p))
+            elif content and content[-1]["type"] == "text":
+                content[-1]["text"] += p
+            else:
+                content.append({"type": "text", "text": p})
+        messages = [{"role": "user", "content": content}]
+        try:
+            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
+                                                      enable_thinking=think)
+        except TypeError:
+            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        return text, images
+
+    @staticmethod
+    def _inputs(text, images):
+        inp = {"prompt": text}
+        if images:
+            inp["multi_modal_data"] = {"image": images}
+        return inp
+
+    # ------------------------------------------------------------ requests
+    def submit_generate(self, req: GenRequest) -> Future:
+        from vllm import SamplingParams
+        think = req.think or self.think
+        text, images = self._prompt(req.parts, think)
+        sp = SamplingParams(max_tokens=req.max_new_tokens, temperature=req.temperature,
+                            top_p=req.top_p if req.temperature > 0 else 1.0,
+                            repetition_penalty=req.repetition_penalty or 1.0,
+                            logprobs=0 if req.logprobs else None)
+
+        def post(out):
+            o = out.outputs[0]
+            raw = o.text
+            ids = list(o.token_ids)
+            if ids and ids[-1] in self.eos_ids:  # like the other backends: the stop token is not part of the answer
+                ids = ids[:-1]
+            lps = None
+            if req.logprobs and o.logprobs is not None:
+                lps = [step[tid].logprob for tid, step in zip(ids, o.logprobs)]
+            text_ = strip_thinking(raw) if (think or "</think>" in raw) else raw.strip()
+            return GenResult(text=text_, raw_text=raw, token_logprobs=lps, num_tokens=len(ids))
+
+        return self._submit(self._inputs(text, images), sp, post)
+
+    def submit_score(self, req: ScoreRequest) -> Future:
+        from vllm import SamplingParams
+        text, images = self._prompt(req.parts, think=False)
+        n = len(self.tokenizer(req.target, add_special_tokens=False)["input_ids"])
+        if n == 0:
+            f: Future = Future()
+            f.set_result([])
+            return f
+        sp = SamplingParams(max_tokens=1, temperature=0.0, prompt_logprobs=0)
+
+        def post(out):
+            ids = out.prompt_token_ids[-n:]
+            plp = out.prompt_logprobs[-n:]
+            return [step[tid].logprob for tid, step in zip(ids, plp)]
+
+        return self._submit(self._inputs(text + req.target, images), sp, post)
+
+    def _submit(self, inputs, params, post) -> Future:
+        fut: Future = Future()
+        self._queue.put((inputs, params, post, fut))
+        return fut
+
+    def close(self):
+        self._queue.put(None)
+
+    # --------------------------------------------------------------- engine
+    def _loop(self):
+        while True:
+            block = not self._live
+            try:
+                item = self._queue.get(timeout=0.5) if block else self._queue.get_nowait()
+                while True:
+                    if item is None:
+                        return
+                    inputs, params, post, fut = item
+                    if not fut.cancelled():
+                        rid = str(next(self._ids))
+                        try:
+                            self.engine.add_request(rid, inputs, params)
+                            self._live[rid] = (post, fut)
+                        except Exception as e:
+                            fut.set_exception(e)
+                    item = self._queue.get_nowait()
+            except queue.Empty:
+                pass
+            if not self._live:
+                continue
+            for rid, (post, fut) in list(self._live.items()):
+                if fut.cancelled():
+                    try:
+                        self.engine.abort_request([rid])
+                    except Exception:
+                        pass
+                    self._live.pop(rid, None)
+            try:
+                outputs = self.engine.step()
+            except Exception as e:
+                tb = traceback.format_exc()
+                for post, fut in self._live.values():
+                    if not fut.done():
+                        fut.set_exception(RuntimeError(f"{e}\n{tb}"))
+                self._live.clear()
+                continue
+            for out in outputs:
+                if not out.finished or out.request_id not in self._live:
+                    continue
+                post, fut = self._live.pop(out.request_id)
+                if fut.done():
+                    continue
+                try:
+                    fut.set_result(post(out))
+                except Exception as e:
+                    fut.set_exception(e)

simit/config.py

@@ -0,0 +1,112 @@
+"""SIMIT hyperparameters. Defaults follow the paper; ``SIMIT.tune`` fits the
+adaptive-budget (ABA) and difficulty-filter (DF) parameters on labeled
+validation data."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+from typing import Optional
+
+
+@dataclass
+class SIMITConfig:
+    # --- budget -----------------------------------------------------------
+    #: Maximum number of imagined demonstrations per query (K_max).
+    k_max: int = 4
+    #: Adaptive budget allocation: fewer demos for confidently answered queries.
+    use_aba: bool = True
+    #: ABA parameters (target error rate and log-accuracy curve A(K) = A1 + B ln K).
+    epsilon: float = 0.14
+    A0: float = 0.35
+    A1: float = 0.60
+    B: float = 0.2
+    #: Difficulty filtering: keep demos whose answer confidence lies in [t_low, t_high].
+    use_df: bool = True
+    t_low: float = 0.2
+    t_high: float = 0.9
+
+    # --- synthesis ----------------------------------------------------------
+    # ``None`` means "the backend's default" (e.g. Lance: decomposed synthesis,
+    # natural images only, no critic -- see ``Backend.pipeline_defaults``).
+    #: Triplet synthesis: "batch" (one call for K triplets, the paper's prompt)
+    #: or "decomposed" (question, answer, description in three short calls).
+    synthesis: Optional[str] = None
+    #: Diversity-encouraging triplet prompt (the paper's "DP").
+    diversity_prompt: bool = True
+    #: Use the structured skill library (False: realize everything as a natural image).
+    use_skills: Optional[bool] = None
+    #: Critic verification of every realized image.
+    verify: Optional[bool] = None
+    #: Minimum critic score (0-100) for an image to be accepted.
+    verify_threshold: int = 50
+    #: Run the critic in thinking mode (None: backend default; BAGEL: True).
+    verify_think: Optional[bool] = None
+    verify_max_new_tokens: int = 300
+    #: Realization attempts per demonstration slot (breadth-first slot filling).
+    attempts_per_slot: int = 4
+    #: Repair/regeneration rounds a failed candidate gets in the retry phase.
+    repair_retries: int = 2
+    verify_rounds: int = 5
+    #: Side length of natively generated images.
+    image_size: int = 400
+    #: Start every remaining candidate attempt at once (up to K_max) instead of only as many as the
+    #: budget still needs. Lower latency for a lone query on an otherwise idle GPU, wasted work under
+    #: load. None: on for ``imagine``, off for ``imagine_batch`` and ``tune``.
+    speculative: Optional[bool] = None
+
+    # --- answering ------------------------------------------------------------
+    answer_max_new_tokens: int = 128
+    #: Random seed for image generation (None: random).
+    seed: Optional[int] = None
+
+    def thresholds(self) -> dict[int, float]:
+        """ABA confidence thresholds tau_K (Appendix D, Algorithm 1)."""
+        return compute_thresholds(self.epsilon, self.A0, self.A1, self.B, self.k_max)
+
+    def budget(self, p0: Optional[float]) -> int:
+        """Number of demonstrations K*(q) for zero-shot confidence ``p0``."""
+        if not self.use_aba or p0 is None:
+            return self.k_max
+        tau = self.thresholds()
+        for k in range(self.k_max + 1):
+            if p0 >= tau[k]:
+                return k
+        return self.k_max
+
+    def in_band(self, c: Optional[float]) -> bool:
+        return (not self.use_df) or (c is not None and self.t_low <= c <= self.t_high)
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    def save(self, path) -> None:
+        Path(path).write_text(json.dumps(self.to_dict(), indent=2))
+
+    @classmethod
+    def load(cls, path) -> "SIMITConfig":
+        data = json.loads(Path(path).read_text())
+        names = {f.name for f in fields(cls)}
+        return cls(**{k: v for k, v in data.items() if k in names})
+
+    def replace(self, **kw) -> "SIMITConfig":
+        data = self.to_dict()
+        data.update(kw)
+        return SIMITConfig(**data)
+
+
+def compute_thresholds(epsilon: float, A0: float, A1: float, B: float, k_max: int, eta: float = 1e-3) -> dict[int, float]:
+    import math
+
+    def clip(a):
+        return min(max(a, eta), 1 - eta)
+
+    a0 = clip(A0)
+    o0 = a0 / (1 - a0)
+    tau = {0: 1 - epsilon}
+    for k in range(1, k_max + 1):
+        ak = clip(A1 + B * math.log(k))
+        ok = ak / (1 - ak)
+        tau[k] = (1 - epsilon) * o0 / (epsilon * ok + (1 - epsilon) * o0)
+    return tau

simit/metrics.py

@@ -0,0 +1,113 @@
+"""Answer-scoring metrics for ``SIMIT.tune`` (higher is better, per sample in [0, 1]).
+
+Every metric is ``metric(prediction: str, references: list[str]) -> float``;
+pass any such callable to ``tune(metric=...)``.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Callable, Sequence, Union
+
+_ARTICLES = {"a", "an", "the"}
+_PUNCT = re.compile(r"[^\w\s]")
+_NUMBER_WORDS = {"none": "0", "zero": "0", "one": "1", "two": "2", "three": "3", "four": "4", "five": "5",
+                 "six": "6", "seven": "7", "eight": "8", "nine": "9", "ten": "10"}
+
+
+def normalize(text: str) -> str:
+    """Lowercase, drop punctuation and articles, map number words to digits."""
+    text = _PUNCT.sub(" ", str(text).lower())
+    words = [_NUMBER_WORDS.get(w, w) for w in text.split() if w not in _ARTICLES]
+    return " ".join(words)
+
+
+def _refs(references) -> list[str]:
+    if isinstance(references, str):
+        return [references]
+    return [str(r) for r in references]
+
+
+def exact_match(prediction: str, references) -> float:
+    p = normalize(prediction)
+    return float(any(p == normalize(r) for r in _refs(references)))
+
+
+def vqa_accuracy(prediction: str, references) -> float:
+    """VQAv2-style soft accuracy: min(1, #matching annotators / 3)."""
+    p = normalize(prediction)
+    refs = _refs(references)
+    hits = sum(p == normalize(r) for r in refs)
+    return min(1.0, hits / 3.0) if len(refs) > 1 else float(hits > 0)
+
+
+def contains(prediction: str, references) -> float:
+    p = normalize(prediction)
+    return float(any(normalize(r) and normalize(r) in p for r in _refs(references)))
+
+
+def _levenshtein(a: str, b: str) -> int:
+    if len(a) < len(b):
+        a, b = b, a
+    prev = list(range(len(b) + 1))
+    for i, ca in enumerate(a, 1):
+        cur = [i]
+        for j, cb in enumerate(b, 1):
+            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
+        prev = cur
+    return prev[-1]
+
+
+def anls(prediction: str, references, threshold: float = 0.5) -> float:
+    """Average normalized Levenshtein similarity (DocVQA / InfographicVQA)."""
+    p = " ".join(str(prediction).lower().split())
+    best = 0.0
+    for r in _refs(references):
+        r = " ".join(r.lower().split())
+        d = _levenshtein(p, r) / max(len(p), len(r), 1)
+        best = max(best, 1 - d if d < threshold else 0.0)
+    return best
+
+
+def relaxed_accuracy(prediction: str, references, tolerance: float = 0.05) -> float:
+    """ChartQA: numbers within 5% relative error, otherwise exact match."""
+    def num(s):
+        s = str(s).strip().rstrip("%").replace(",", "")
+        try:
+            return float(s)
+        except ValueError:
+            return None
+
+    pv = num(prediction)
+    for r in _refs(references):
+        rv = num(r)
+        if pv is not None and rv is not None:
+            if (rv == 0 and pv == 0) or (rv != 0 and abs(pv - rv) / abs(rv) <= tolerance):
+                return 1.0
+        elif normalize(prediction) == normalize(r):
+            return 1.0
+    return 0.0
+
+
+_LETTER = re.compile(r"\b([A-J])\b")
+
+
+def multiple_choice(prediction: str, references) -> float:
+    """Compare the first option letter in the prediction with the reference letter."""
+    m = _LETTER.search(str(prediction).strip().upper()) or re.match(r"\s*([A-J])", str(prediction).upper())
+    letter = m.group(1) if m else str(prediction).strip()[:1].upper()
+    return float(any(letter == str(r).strip().upper()[:1] for r in _refs(references)))
+
+
+METRICS: dict[str, Callable] = {
+    "exact_match": exact_match, "vqa_accuracy": vqa_accuracy, "contains": contains, "anls": anls,
+    "relaxed_accuracy": relaxed_accuracy, "multiple_choice": multiple_choice,
+}
+
+
+def get_metric(metric: Union[str, Callable]) -> Callable[[str, Sequence[str]], float]:
+    if callable(metric):
+        return metric
+    if metric not in METRICS:
+        raise ValueError(f"unknown metric {metric!r}; choose from {sorted(METRICS)} or pass a callable")
+    return METRICS[metric]