simit 0.1.0__py3-none-any.whl

simit/__init__.py

@@ -0,0 +1,39 @@
+"""SIMIT-ICL: self-improving vision-language models via imagination at test time.
+
+    from simit import SIMIT
+    model = SIMIT.from_pretrained("ByteDance-Seed/BAGEL-7B-MoT")
+    demos = model.imagine(image, question)
+    answer = model.answer(image, question, demos)
+"""
+
+__version__ = "0.1.0"
+
+# Heavy modules (torch, transformers) load lazily so that light imports such
+# as ``from simit import Skill`` -- and the render worker processes -- stay fast.
+_LAZY = {
+    "SIMIT": ("simit.api", "SIMIT"),
+    "SIMITConfig": ("simit.config", "SIMITConfig"),
+    "Demo": ("simit.types", "Demo"),
+    "Skill": ("simit.skills.base", "Skill"),
+    "SpecError": ("simit.skills.base", "SpecError"),
+    "Realization": ("simit.skills.base", "Realization"),
+    "extract_fenced_block": ("simit.skills.base", "extract_fenced_block"),
+    "save_demos": ("simit.types", "save_demos"),
+    "load_demos": ("simit.types", "load_demos"),
+    "default_skills": ("simit.skills.builtin", "default_skills"),
+    "metrics": ("simit.metrics", None),
+}
+
+__all__ = list(_LAZY) + ["__version__"]
+
+
+def __getattr__(name):
+    if name not in _LAZY:
+        raise AttributeError(f"module 'simit' has no attribute {name!r}")
+    import importlib
+
+    module_name, attr = _LAZY[name]
+    module = importlib.import_module(module_name)
+    value = module if attr is None else getattr(module, attr)
+    globals()[name] = value
+    return value

simit/api.py

@@ -0,0 +1,258 @@
+"""The user-facing ``SIMIT`` class."""
+
+from __future__ import annotations
+
+import logging
+from collections import OrderedDict
+from typing import Callable, Iterable, Optional, Sequence, Union
+
+from PIL import Image
+
+from .backends import load_backend, load_image_generator, registered_backend
+from .backends.base import Backend
+from .config import SIMITConfig
+from .pipeline.core import Pipeline, Runtime
+from .skills.base import Skill
+from .skills.builtin import default_skills
+from .skills.pool import RenderPool
+from .types import Demo, Imagination, ZeroShot
+from .utils import ImageLike, image_hash, to_pil
+
+log = logging.getLogger("simit")
+
+
+class SIMIT:
+    """A vision-language model that imagines its own in-context demonstrations.
+
+    >>> model = SIMIT.from_pretrained("ByteDance-Seed/BAGEL-7B-MoT")
+    >>> demos = model.imagine(image, question)
+    >>> answer = model.answer(image, question, demos)
+    """
+
+    def __init__(self, backend: Backend, *, image_generator=None, skills: Optional[Sequence[Skill]] = None,
+                 config: Optional[SIMITConfig] = None, render_workers: Optional[int] = None,
+                 model_id: str = "", verbose: bool = False):
+        if verbose:
+            logging.basicConfig(level=logging.INFO, format="%(asctime)s [simit] %(message)s")
+            log.setLevel(logging.INFO)
+        self.backend = backend
+        self.image_generator = image_generator
+        self.model_id = model_id or getattr(backend, "model_id", backend.name)
+        self._runtime = Runtime()
+        self._pool = RenderPool(render_workers)
+        self._pipe = Pipeline(backend, skills if skills is not None else default_skills(),
+                              config or SIMITConfig(), image_generator=image_generator, pool=self._pool,
+                              runtime=self._runtime)
+        self._zs_cache: "OrderedDict[tuple, ZeroShot]" = OrderedDict()
+
+    # ------------------------------------------------------------ loading
+    @classmethod
+    def from_pretrained(cls, model: Union[str, object], *, processor=None,
+                        image_generator: Union[str, object, None] = None, skills: Optional[Sequence[Skill]] = None,
+                        config: Optional[SIMITConfig] = None, device: Optional[str] = None,
+                        image_generator_device: Optional[str] = None, render_workers: Optional[int] = None,
+                        verbose: bool = False, **backend_kwargs) -> "SIMIT":
+        """Load a model by Hugging Face id / local path, or wrap an already
+        loaded ``transformers`` model (pass its ``processor`` too).
+
+        ``image_generator``: a diffusers model id (e.g.
+        ``"black-forest-labs/FLUX.2-klein-4B"``), a diffusers pipeline, or a
+        callable ``(prompt, width, height) -> PIL.Image``. It is used for
+        natural images when the model cannot generate images itself; omit it
+        to use the structured skills only."""
+        gen = None
+        gen_device = image_generator_device or _second_device(device)
+        if image_generator is not None and not isinstance(model, Backend) and registered_backend(model) is None:
+            # Load the tool first so the VLM (vLLM / device_map="auto") sizes itself around it.
+            gen = load_image_generator(image_generator, device=gen_device)
+        backend = load_backend(model, processor=processor, device=device, **backend_kwargs)
+        if image_generator is not None and backend.native_image_generation:
+            log.info("%s generates images natively; ignoring image_generator", backend.name)
+        elif image_generator is not None and gen is None:
+            gen = load_image_generator(image_generator, device=gen_device)
+        model_id = model if isinstance(model, str) else getattr(getattr(model, "config", None), "_name_or_path", "")
+        return cls(backend, image_generator=gen, skills=skills, config=config, render_workers=render_workers,
+                   model_id=model_id, verbose=verbose)
+
+    # --------------------------------------------------------- properties
+    @property
+    def config(self) -> SIMITConfig:
+        return self._pipe.config
+
+    @config.setter
+    def config(self, cfg: SIMITConfig):
+        self._pipe.config = cfg
+
+    @property
+    def skills(self) -> list[Skill]:
+        return list(self._pipe.skills)
+
+    def add_skill(self, skill: Skill, replace: bool = False) -> None:
+        """Register a custom skill (it joins the router's categories)."""
+        if not isinstance(skill, Skill):
+            raise TypeError("expected a simit.Skill instance")
+        current = [s for s in self._pipe.skills if not (replace and s.name == skill.name)]
+        self._pipe.set_skills(current + [skill])
+
+    def remove_skill(self, name: str) -> None:
+        self._pipe.set_skills([s for s in self._pipe.skills if s.name != name])
+
+    # ------------------------------------------------------------ helpers
+    def _run(self, coro):
+        return self._runtime.run(coro)
+
+    def _mnt(self, max_new_tokens):
+        return max_new_tokens or self.config.answer_max_new_tokens
+
+    async def _zero_shot(self, image: Image.Image, question: str, mnt: int) -> ZeroShot:
+        key = (image_hash(image), question, mnt)
+        if key in self._zs_cache:
+            self._zs_cache.move_to_end(key)
+            return self._zs_cache[key]
+        zs = await self._pipe.zero_shot(image, question, mnt)
+        self._zs_cache[key] = zs
+        while len(self._zs_cache) > 4096:
+            self._zs_cache.popitem(last=False)
+        return zs
+
+    async def _imagine(self, image, question, k, fill, mnt, alone: bool, time_limit=None, on_demo=None,
+                       on_progress=None):
+        image = to_pil(image)
+        zs = await self._zero_shot(image, question, mnt)
+        spec = self.config.speculative
+        if on_progress is not None:
+            self._pipe.on_call = on_progress
+        try:
+            return await self._pipe.imagine(image, question, k=k, fill=fill, max_new_tokens=mnt, zero_shot=zs,
+                                            speculative=alone if spec is None else spec, time_limit=time_limit,
+                                            on_demo=on_demo)
+        finally:
+            if on_progress is not None:
+                self._pipe.on_call = None
+
+    async def _answer(self, image, question, demos, mnt):
+        image = to_pil(image)
+        if not demos:
+            return (await self._zero_shot(image, question, mnt)).answer
+        return await self._pipe.answer(image, question, demos, mnt)
+
+    # -------------------------------------------------------------- public
+    def imagine(self, image: ImageLike, question: str, k: Optional[int] = None, *,
+                return_details: bool = False, max_new_tokens: Optional[int] = None,
+                time_limit: Optional[float] = None, on_demo: Optional[Callable[[Demo], None]] = None,
+                on_progress: Optional[Callable[[str], None]] = None):
+        """Imagine demonstrations similar to the query ``(image, question)``.
+
+        By default the number of demos is chosen adaptively from the model's
+        zero-shot confidence (0 for queries it is already sure about) and
+        candidates outside the confidence band are filtered out. ``k=n``
+        forces exactly ``n`` verified demos. ``return_details=True`` returns an
+        :class:`Imagination` (zero-shot answer, budget, all candidates, stats).
+        A single call runs its candidate attempts in parallel (see
+        ``SIMITConfig.speculative``); use :meth:`imagine_batch` for many queries.
+
+        ``time_limit`` (seconds, not counting the zero-shot pass) stops imagining
+        and returns the demos accepted so far; ``on_demo(demo)`` is called as
+        each demo is accepted and ``on_progress(stage)`` as each model call
+        starts ("synthesis", "route", "image", "render:<skill>", "critic", ...),
+        e.g. to show progress in a UI."""
+        res: Imagination = self._run(self._imagine(image, question, k, False, self._mnt(max_new_tokens), True,
+                                                   time_limit, on_demo, on_progress))
+        return res if return_details else res.demos
+
+    def answer(self, image: ImageLike, question: str, demos: Optional[Sequence[Demo]] = None, *,
+               max_new_tokens: Optional[int] = None) -> str:
+        """Answer the query using the demonstrations in context (no demos:
+        the zero-shot answer)."""
+        return self._run(self._answer(image, question, list(demos or []), self._mnt(max_new_tokens)))
+
+    def greedy(self, image: ImageLike, question: str, *, max_new_tokens: Optional[int] = None) -> str:
+        """The standard zero-shot greedy answer (cached, shared with ``imagine``)."""
+        return self.zero_shot(image, question, max_new_tokens=max_new_tokens).answer
+
+    def zero_shot(self, image: ImageLike, question: str, *, max_new_tokens: Optional[int] = None) -> ZeroShot:
+        """The greedy answer with its confidence ``p0`` (what the adaptive budget
+        uses: ``model.config.budget(zs.confidence)`` demos). Cached and reused by
+        a following ``imagine`` of the same query."""
+        return self._run(self._zero_shot(to_pil(image), question, self._mnt(max_new_tokens)))
+
+    def __call__(self, image: ImageLike, question: str, k: Optional[int] = None, **kw) -> str:
+        """Imagine, then answer with the imagined demonstrations."""
+        return self.answer(image, question, self.imagine(image, question, k, **kw), **kw)
+
+    def imagine_batch(self, queries: Iterable, k: Optional[int] = None, *, return_details: bool = False,
+                      max_new_tokens: Optional[int] = None, concurrency: int = 16) -> list:
+        """``imagine`` for many ``(image, question)`` pairs at once; their model
+        calls are batched together (much faster than a loop)."""
+        import asyncio
+        queries = [(q["image"], q["question"]) if isinstance(q, dict) else tuple(q) for q in queries]
+        mnt = self._mnt(max_new_tokens)
+
+        async def run():
+            sem = asyncio.Semaphore(concurrency)
+
+            async def one(img, q):
+                async with sem:
+                    return await self._imagine(img, q, k, False, mnt, False)
+            return await asyncio.gather(*[one(img, q) for img, q in queries])
+
+        res = self._run(run())
+        return res if return_details else [r.demos for r in res]
+
+    def answer_batch(self, queries: Iterable, demos: Optional[Sequence[Sequence[Demo]]] = None, *,
+                     max_new_tokens: Optional[int] = None) -> list[str]:
+        import asyncio
+        queries = [(q["image"], q["question"]) if isinstance(q, dict) else tuple(q) for q in queries]
+        demos = demos if demos is not None else [[] for _ in queries]
+        mnt = self._mnt(max_new_tokens)
+
+        async def run():
+            return await asyncio.gather(*[self._answer(img, q, list(d), mnt) for (img, q), d in zip(queries, demos)])
+        return self._run(run())
+
+    def tune(self, val_set, metric: Union[str, Callable] = "exact_match", n_trials: int = 100, *,
+             k_max_choices: Optional[Sequence[int]] = None, cache_dir=None, seed: int = 0,
+             concurrency: int = 16, max_new_tokens: Optional[int] = None, apply: bool = True,
+             max_candidates: Optional[float] = None, max_runtime_ratio: Optional[float] = None):
+        """Fit ABA + DF hyperparameters on labeled validation data.
+
+        ``val_set``: iterable of ``(image, question, answer_or_answers)`` or
+        dicts with ``image``/``question``/``answer(s)`` (and optionally their own
+        ``metric``, for validation sets that mix benchmarks). ``metric``: a name from
+        :mod:`simit.metrics` or ``fn(prediction, references) -> float``.
+        Synthesis runs once per query (cached, optionally on disk under
+        ``cache_dir``); each Optuna trial only re-selects and re-answers.
+
+        ``max_candidates`` caps the mean number of candidates synthesized per
+        query (the dominant runtime cost), e.g. ``1.0`` for roughly a quarter
+        of K_max=4 as in the paper. ``max_runtime_ratio`` (e.g. ``25``) caps the
+        estimated batched runtime relative to zero-shot instead, from timings
+        measured on this machine while tuning. ``None`` for both optimizes
+        accuracy alone."""
+        from .tune import run_tune
+        return run_tune(self, val_set, metric=metric, n_trials=n_trials, k_max_choices=k_max_choices,
+                        cache_dir=cache_dir, seed=seed, concurrency=concurrency, max_new_tokens=max_new_tokens,
+                        apply=apply, max_candidates=max_candidates, max_runtime_ratio=max_runtime_ratio)
+
+    def close(self):
+        self._pool.close()
+        self.backend.close()
+        if self.image_generator is not None:
+            self.image_generator.close()
+        self._runtime.close()
+
+    def __repr__(self):
+        return (f"SIMIT(model={self.model_id!r}, backend={self.backend.name}, skills={len(self.skills)}, "
+                f"image_generator={'native' if self.backend.native_image_generation else type(self.image_generator).__name__ if self.image_generator else None})")
+
+
+def _second_device(device: Optional[str]) -> Optional[str]:
+    """Put a separate image generator on another GPU when there is one."""
+    try:
+        import torch
+        n = torch.cuda.device_count()
+    except Exception:
+        return None
+    if n >= 2:
+        return f"cuda:{n - 1}"
+    return device

simit/backends/__init__.py

@@ -0,0 +1,112 @@
+"""Model backends. ``load_backend`` picks one from a model id/path/object:
+
+* BAGEL checkpoints -> :class:`~simit.backends.bagel.BagelBackend` (native image generation)
+* Lance checkpoints -> :class:`~simit.backends.lance.LanceBackend` (native image generation)
+* anything else     -> :class:`~simit.backends.hf.HFBackend` (any transformers VLM)
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Optional
+
+from .base import Backend, GenRequest, GenResult, ImageRequest, ScoreRequest
+
+__all__ = ["Backend", "GenRequest", "GenResult", "ScoreRequest", "ImageRequest", "load_backend",
+           "load_image_generator", "register_backend", "registered_backend"]
+
+# name -> (detector(model_id, config) -> bool, factory(model, **kw) -> Backend)
+_REGISTRY: dict = {}
+
+
+def register_backend(name: str, detect, factory):
+    """Add a backend for a new model family (checked before the generic HF backend)."""
+    _REGISTRY[name] = (detect, factory)
+
+
+def _read_config(model_id: str) -> dict:
+    path = Path(model_id)
+    try:
+        if path.exists():
+            f = path / "config.json"
+            return json.loads(f.read_text()) if f.exists() else {}
+        from huggingface_hub import hf_hub_download
+        return json.loads(Path(hf_hub_download(model_id, "config.json")).read_text())
+    except Exception:
+        return {}
+
+
+def _is_bagel(model_id: str, cfg: dict) -> bool:
+    return cfg.get("model_type") == "bagel" or "BagelForConditionalGeneration" in cfg.get("architectures", [])
+
+
+def _is_lance(model_id: str, cfg: dict) -> bool:
+    return cfg.get("model_name") == "Lance" or Path(model_id).name.lower() == "lance"
+
+
+def _bagel(model, **kw):
+    from .bagel import BagelBackend
+    return BagelBackend(model, **kw)
+
+
+def _lance(model, **kw):
+    from .lance import LanceBackend
+    return LanceBackend(model, **kw)
+
+
+register_backend("bagel", _is_bagel, _bagel)
+register_backend("lance", _is_lance, _lance)
+
+
+def registered_backend(model) -> Optional[str]:
+    """Name of the registered model-family backend (e.g. ``"bagel"``) for a
+    model id/path, or None for generic transformers/vLLM models."""
+    if not isinstance(model, str):
+        return None
+    cfg = _read_config(model)
+    return next((name for name, (detect, _) in _REGISTRY.items() if detect(model, cfg)), None)
+
+
+def _vllm_available() -> bool:
+    import importlib.util
+    return importlib.util.find_spec("vllm") is not None
+
+
+def load_backend(model, processor=None, device: Optional[str] = None, engine: str = "auto", **kw) -> Backend:
+    """``engine`` (standard VLMs only): ``"transformers"``, ``"vllm"``, or
+    ``"auto"`` (vLLM when installed, otherwise transformers)."""
+    if isinstance(model, Backend):
+        return model
+    if not isinstance(model, str):  # an already-loaded transformers model
+        from .hf import HFBackend
+        return HFBackend(model, processor=processor, **kw)
+    cfg = _read_config(model)
+    for name, (detect, factory) in _REGISTRY.items():
+        if detect(model, cfg):
+            if device is not None:
+                kw["device"] = device
+            return factory(model, **kw)
+    if engine not in ("auto", "transformers", "vllm"):
+        raise ValueError(f"engine must be 'auto', 'transformers' or 'vllm', got {engine!r}")
+    if engine == "vllm" or (engine == "auto" and _vllm_available() and processor is None):
+        from .vllm import VLLMBackend
+        return VLLMBackend(model, device=device, **kw)
+    from .hf import HFBackend
+    if device is not None:
+        kw.setdefault("device_map", device)
+    return HFBackend(model, processor=processor, **kw)
+
+
+def load_image_generator(gen, device: Optional[str] = None):
+    """A text-to-image tool for models without native generation."""
+    if gen is None or isinstance(gen, Backend):
+        return gen
+    from .imagegen import CallableImageGenerator, DiffusersImageGenerator
+    if isinstance(gen, str):
+        return DiffusersImageGenerator(gen, device=device)
+    if hasattr(gen, "__call__") and hasattr(gen, "components"):  # a diffusers pipeline object
+        return DiffusersImageGenerator(gen, device=device)
+    if callable(gen):
+        return CallableImageGenerator(gen)
+    raise TypeError(f"unsupported image_generator: {type(gen).__name__}")