simit 0.1.0__py3-none-any.whl

simit/backends/bagel.py

@@ -0,0 +1,359 @@
+"""BAGEL-7B-MoT (ByteDance-Seed/BAGEL-7B-MoT) on the batched MoT engine.
+
+Prompt formats reproduce the SIMIT research setup: ``chat`` requests use the
+serving template the synthesis pipeline was developed with (system prompt,
+user/assistant turns, repetition penalty 1.05); ``answer`` requests use
+BAGEL's interleaved inference format (each text part its own turn), which
+the evaluation and confidence scores were computed with.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+from PIL import Image
+
+from ..utils import image_hash, strip_thinking, to_pil
+from .base import Backend, GenRequest, GenResult, ImageRequest, ScoreRequest, chained
+from .mot.engine import ImageChunk, ImageJob, MoTEngine, TextChunk, TextJob, auto_kv_tokens
+from .mot.modeling import (DeviceMixin, FluxAutoEncoder, MLPConnector, MoTLanguageModel, PositionTable, SiglipNaViT,
+                           TimestepEmbedder)
+
+DEFAULT_SYSTEM_PROMPT = "You are BAGEL, a helpful assistant created by ByteDance."
+CHAT_THINK_PROMPT = ("You should first think about the reasoning process in the mind "
+                     "and then provide the user with the answer.")
+ANSWER_THINK_PROMPT = ("You should first think about the reasoning process in the mind and then provide the "
+                       "user with the answer. \nThe reasoning process is enclosed within <think> </think> tags, "
+                       "i.e. <think> reasoning process here </think> answer here")
+
+
+# -------------------------------------------------------------- image prep
+def _resize(img: Image.Image, max_size: int, min_size: int, stride: int, max_pixels: int) -> Image.Image:
+    """BAGEL's MaxLongEdgeMinShortEdgeResize (bicubic, antialiased)."""
+    def divisible(v):
+        return max(stride, int(round(v / stride) * stride))
+
+    def scaled(w, h, s):
+        return divisible(round(w * s)), divisible(round(h * s))
+
+    w, h = img.size
+    s = min(max_size / max(w, h), 1.0)
+    s = max(s, min_size / min(w, h))
+    nw, nh = scaled(w, h, s)
+    if nw * nh > max_pixels:
+        nw, nh = scaled(nw, nh, max_pixels / (nw * nh))
+    if max(nw, nh) > max_size:
+        nw, nh = scaled(nw, nh, max_size / max(nw, nh))
+    if (nw, nh) == (w, h):
+        return img
+    from torchvision.transforms import InterpolationMode
+    from torchvision.transforms import functional as TF
+    return TF.resize(img, [nh, nw], InterpolationMode.BICUBIC, antialias=True)
+
+
+def _patchify(img: Image.Image, patch: int, max_side: int):
+    from torchvision.transforms import functional as TF
+    t = TF.to_tensor(img)
+    t = (t - 0.5) / 0.5
+    c, h, w = t.shape
+    hp, wp = h // patch, w // patch
+    patches = t.reshape(c, hp, patch, wp, patch).permute(1, 3, 2, 4, 0).reshape(hp * wp, patch * patch * c)
+    pos = (torch.arange(hp)[:, None] * max_side + torch.arange(wp)).flatten()
+    return patches, pos
+
+
+class BagelModel(DeviceMixin):
+    """Weights + modality-specific pieces the MoT engine calls into."""
+
+    latent_patch = 2
+    latent_channels = 16
+    vae_downsample = 8
+    num_timesteps = 50
+    timestep_shift = 3.0
+    cfg_text_scale = 4.0
+    cfg_interval = (0.4, 1.0)
+    cfg_renorm_min = 0.0
+
+    @classmethod
+    def load(cls, model_id: str = "ByteDance-Seed/BAGEL-7B-MoT", device: str = "cuda",
+             dtype=torch.bfloat16) -> "BagelModel":
+        """From a Hub id (downloaded on first use) or a local path."""
+        path = Path(model_id)
+        if not path.exists():
+            from huggingface_hub import snapshot_download
+            path = Path(snapshot_download(model_id, allow_patterns=[
+                "*.json", "ema.safetensors", "ae.safetensors", "*.txt", "tokenizer*", "vocab*", "merges*"]))
+        model = cls(str(path), device=device, dtype=dtype)
+        model.model_id = model_id
+        return model
+
+    def __init__(self, path: str, device: str = "cuda", dtype=torch.bfloat16):
+        self.model_id = str(path)
+        self.device = torch.device(device)
+        self.dtype = dtype
+        path = Path(path)
+        cfg = json.loads((path / "config.json").read_text())
+        llm_cfg = json.loads((path / "llm_config.json").read_text())
+        vit_cfg = json.loads((path / "vit_config.json").read_text())
+        self.vit_patch = vit_cfg["patch_size"]
+        self.vit_max_side = cfg.get("vit_max_num_patch_per_side", 70)
+        self.latent_downsample = self.vae_downsample * self.latent_patch
+        self.latent_dim = self.latent_patch ** 2 * self.latent_channels
+        self.max_latent_side = 64
+
+        hidden = llm_cfg["hidden_size"]
+        with torch.device("meta"):
+            self.lm = MoTLanguageModel(llm_cfg)
+            self.vit = SiglipNaViT(vit_cfg, vit_cfg["num_hidden_layers"] - 1)
+            self.connector = MLPConnector(vit_cfg["hidden_size"], hidden)
+            self.vit_pos_embed = PositionTable(self.vit_max_side, hidden)
+            self.latent_pos_embed = PositionTable(self.max_latent_side, hidden)
+            self.time_embedder = TimestepEmbedder(hidden)
+            self.vae2llm = torch.nn.Linear(self.latent_dim, hidden)
+            self.llm2vae = torch.nn.Linear(hidden, self.latent_dim)
+        modules = {
+            "language_model": self.lm, "vit_model": self.vit, "connector": self.connector,
+            "vit_pos_embed": self.vit_pos_embed, "latent_pos_embed": self.latent_pos_embed,
+            "time_embedder": self.time_embedder, "vae2llm": self.vae2llm, "llm2vae": self.llm2vae,
+        }
+        self._load(path / "ema.safetensors", modules)
+        self.lm.rotary = type(self.lm.rotary)(hidden // llm_cfg["num_attention_heads"],
+                                              llm_cfg.get("rope_theta", 1e6)).to(self.device)
+        with torch.device("meta"):
+            self.vae = FluxAutoEncoder()
+        self._load(path / "ae.safetensors", {"": self.vae}, dtype=torch.float32)
+
+        from transformers import Qwen2TokenizerFast  # AutoTokenizer would warn about model_type "bagel"
+        tok = Qwen2TokenizerFast.from_pretrained(str(path))
+        missing = [t for t in ("<|im_start|>", "<|im_end|>", "<|vision_start|>", "<|vision_end|>")
+                   if t not in tok.get_vocab()]
+        if missing:
+            tok.add_tokens(missing)
+        self.tokenizer = tok
+        self.bos = tok.convert_tokens_to_ids("<|im_start|>")
+        self.eos = tok.convert_tokens_to_ids("<|im_end|>")
+        self.start_of_image = tok.convert_tokens_to_ids("<|vision_start|>")
+        self.end_of_image = tok.convert_tokens_to_ids("<|vision_end|>")
+        self.stop_ids = (self.eos, tok.convert_tokens_to_ids("<|endoftext|>"))
+
+    def _load(self, file: Path, modules: dict, dtype=None):
+        from safetensors import safe_open
+        dtype = dtype or self.dtype
+        with safe_open(str(file), framework="pt", device=str(self.device)) as f:
+            keys = set(f.keys())
+            for prefix, mod in modules.items():
+                pre = f"{prefix}." if prefix else ""
+                sd = {}
+                for name, _ in list(mod.named_parameters()) + list(mod.named_buffers()):
+                    key = pre + name
+                    if key in keys:
+                        sd[name] = f.get_tensor(key).to(dtype)
+                missing = [n for n, _ in mod.named_parameters() if n not in sd]
+                if missing:
+                    raise RuntimeError(f"{file.name}: missing weights for {prefix}: {missing[:5]}")
+                mod.load_state_dict(sd, strict=False, assign=True)
+                mod.eval()
+
+    # ---------------------------------------------------------- understanding
+    def preprocess_image(self, image: Image.Image) -> ImageChunk:
+        img = _resize(to_pil(image), 1024, 512, 16, 14 * 14 * 9 * 1024)   # BAGEL's VAE-side resize first
+        img = _resize(img, 980, 224, self.vit_patch, 14 * 14 * 9 * 1024)
+        patches, pos = _patchify(img, self.vit_patch, self.vit_max_side)
+        return ImageChunk(n=patches.shape[0] + 2, key=image_hash(img), payload=(patches, pos))
+
+    def encode_images(self, payloads):
+        patches = torch.cat([p for p, _ in payloads]).to(self.device, self.dtype)
+        pos = torch.cat([q for _, q in payloads]).to(self.device)
+        lens = [p.shape[0] for p, _ in payloads]
+        cu = torch.tensor([0] + list(_cumsum(lens)), dtype=torch.int32, device=self.device)
+        x = self.vit(patches, pos, cu, max(lens))
+        x = self.connector(x) + self.vit_pos_embed(pos).to(x.dtype)
+        return list(x.split(lens))
+
+    # 1-D rope: text positions increase per token, an image block shares one.
+    def text_positions(self, start: int, n: int):
+        return list(range(start, start + n))
+
+    def image_positions(self, start: int, chunk):
+        return [start] * chunk.n
+
+    def latent_positions(self, start: int, h: int, w: int):
+        return [start] * (h * w + 2)
+
+    def position_after_image(self, pos: int, chunk) -> int:
+        return pos + 1
+
+    def position_after(self, chunks, n: int) -> int:
+        pos = 0
+        for c in chunks:
+            if n <= 0:
+                break
+            if isinstance(c, TextChunk):
+                take = min(n, len(c.ids))
+                pos += take
+                n -= take
+            else:
+                pos += 1
+                n -= c.n
+        return pos
+
+    # ------------------------------------------------------------- generation
+    def latent_pos_ids(self, h: int, w: int):
+        return (torch.arange(h)[:, None] * self.max_latent_side + torch.arange(w)).flatten()
+
+    def denoise_schedule(self):
+        n, s = self.num_timesteps, self.timestep_shift
+        t = torch.linspace(1, 0, n)
+        t = s * t / (1 + (s - 1) * t)
+        return [(float(t[i]), float(t[i] - t[i + 1])) for i in range(n - 1)]
+
+    def cfg_active(self, t: float) -> bool:
+        lo, hi = self.cfg_interval
+        return lo < t <= hi and self.cfg_text_scale > 1.0
+
+    def embed_latents(self, x, t: float, pos_ids):
+        tt = torch.full((x.shape[0],), t, device=x.device)
+        out = self.vae2llm(x.to(self.dtype)) + self.time_embedder(tt) + self.latent_pos_embed(pos_ids).to(self.dtype)
+        return out.to(self.dtype)
+
+    def latents_to_velocity(self, h):
+        return self.llm2vae(h)
+
+    def apply_cfg(self, v, v_uncond):
+        v, v_uncond = v.float(), v_uncond.float()
+        guided = v_uncond + self.cfg_text_scale * (v - v_uncond)
+        scale = (v.norm() / (guided.norm() + 1e-8)).clamp(min=self.cfg_renorm_min, max=1.0)  # "global" renorm
+        return guided * scale
+
+    def decode_latents(self, items):
+        out = []
+        p, c = self.latent_patch, self.latent_channels
+        for x, h, w in items:
+            lat = x.reshape(1, h, w, p, p, c)
+            lat = torch.einsum("nhwpqc->nchpwq", lat).reshape(1, c, h * p, w * p)
+            img = self.vae.decode(lat.float())
+            img = ((img * 0.5 + 0.5).clamp(0, 1)[0].permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
+            out.append(Image.fromarray(img))
+        return out
+
+
+def _cumsum(xs):
+    total = 0
+    for x in xs:
+        total += x
+        yield total
+
+
+class BagelBackend(Backend):
+    name = "bagel"
+    native_image_generation = True
+    verify_think = True   # the SIMIT critic runs BAGEL in thinking mode
+
+    def __init__(self, model_id: Union[str, "BagelModel"] = "ByteDance-Seed/BAGEL-7B-MoT", device: Optional[str] = None,
+                 kv_cache_tokens: Optional[int] = None, max_step_tokens: int = 16384,
+                 image_steps: Optional[int] = None, use_cuda_graphs: bool = True):
+        """``model_id``: a Hub id / local path, or a :class:`BagelModel` loaded earlier
+        (e.g. on CPU at startup, then moved per request). ``device``: where to put
+        the weights (default ``cuda:0``; for a loaded model, where it already is)."""
+        if isinstance(model_id, BagelModel):
+            self.model = model_id
+            if device is not None and self.model.device != torch.device(device):
+                self.model.to(device)
+        else:
+            self.model = BagelModel.load(model_id, device=device or "cuda:0")
+        self.model_id = self.model.model_id
+        if image_steps is not None:  # fewer denoising steps: faster, slightly lower fidelity
+            self.model.num_timesteps = int(image_steps)
+        if kv_cache_tokens is None:
+            kv_cache_tokens = auto_kv_tokens(self.model)
+        self.engine = MoTEngine(self.model, kv_tokens=kv_cache_tokens, max_step_tokens=max_step_tokens,
+                                use_cuda_graphs=use_cuda_graphs)
+
+    # ------------------------------------------------------------- prompts
+    def _encode(self, text: str) -> list[int]:
+        return self.model.tokenizer.encode(text, add_special_tokens=False)
+
+    def _chunks(self, parts, style: str, think: bool, open_answer: bool = True):
+        m = self.model
+        chunks = []
+        if style == "chat":
+            system = DEFAULT_SYSTEM_PROMPT + (f" {CHAT_THINK_PROMPT}" if think else "")
+            has_image = any(isinstance(p, Image.Image) for p in parts)
+            if not has_image:
+                prompt = "".join(parts)
+                text = f"<|im_start|>{system}<|im_end|><|im_start|>{prompt}<|im_end|><|im_start|>"
+                return [TextChunk(self._encode(text))]
+            buf = f"<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n"
+            after_image = False
+            for p in parts:
+                if isinstance(p, Image.Image):
+                    chunks.append(TextChunk(self._encode(buf)))
+                    chunks.append(m.preprocess_image(p))
+                    buf, after_image = "", True
+                else:
+                    buf += ("\n" if after_image else "") + p
+                    after_image = False
+            buf += "<|im_end|>\n<|im_start|>assistant\n"
+            chunks.append(TextChunk(self._encode(buf)))
+            return [c for c in chunks if not isinstance(c, TextChunk) or c.ids]
+        # "answer": every text part is its own <|im_start|>...<|im_end|> turn
+        if think:
+            chunks.append(TextChunk([m.bos] + self._encode(ANSWER_THINK_PROMPT) + [m.eos]))
+        for p in parts:
+            if isinstance(p, Image.Image):
+                chunks.append(m.preprocess_image(p))
+            else:
+                chunks.append(TextChunk([m.bos] + self._encode(p) + [m.eos]))
+        if open_answer:
+            chunks.append(TextChunk([m.bos]))
+        return chunks
+
+    # -------------------------------------------------------------- requests
+    def submit_generate(self, req: GenRequest):
+        rp = req.repetition_penalty
+        if rp is None:
+            rp = 1.05 if req.style == "chat" else 1.0
+        job = TextJob(chunks=self._chunks(req.parts, req.style, req.think), max_new_tokens=req.max_new_tokens,
+                      temperature=req.temperature, top_p=req.top_p, repetition_penalty=rp,
+                      logprobs=req.logprobs, stop_ids=self.model.stop_ids)
+        tok = self.model.tokenizer
+        if req.on_text is not None:
+            job.on_tokens = lambda ids: req.on_text(tok.decode(ids, skip_special_tokens=True))
+
+        def post(seq):
+            raw = tok.decode(seq.out_ids, skip_special_tokens=False)
+            raw = raw.split("<|im_end|>")[0].replace("<|im_start|>", "")
+            text = strip_thinking(raw) if req.think else raw.strip()
+            return GenResult(text=text, raw_text=raw, token_logprobs=seq.out_lp if req.logprobs else None,
+                             num_tokens=len(seq.out_ids))
+
+        return chained(self.engine.submit(job), post)
+
+    def submit_score(self, req: ScoreRequest):
+        target = self._encode(req.target)
+        if not target:
+            from concurrent.futures import Future
+            f = Future()
+            f.set_result([])
+            return f
+        chunks = self._chunks(req.parts, req.style, think=False, open_answer=False)
+        if req.style == "chat":
+            chunks[-1] = TextChunk(chunks[-1].ids + target)
+        else:
+            chunks.append(TextChunk([self.model.bos] + target))
+        job = TextJob(chunks=chunks, score_len=len(target), stop_ids=self.model.stop_ids)
+        return chained(self.engine.submit(job), lambda seq: seq.score_lp)
+
+    def submit_image(self, req: ImageRequest):
+        w, h = (max(16, int(round(v / 16)) * 16) for v in (req.width, req.height))
+        text = f"<|im_start|>{req.prompt}<|im_end|><|im_start|>"
+        job = ImageJob(chunks=[TextChunk(self._encode(text))], height=h, width=w, seed=req.seed)
+        return self.engine.submit(job)
+
+    def close(self):
+        self.engine.close()

simit/backends/base.py

@@ -0,0 +1,114 @@
+"""Backend interface: every model (UMM or plain VLM) serves three request
+kinds, each submitted asynchronously and batched by the backend itself."""
+
+from __future__ import annotations
+
+from concurrent.futures import Future
+from dataclasses import dataclass, field
+from typing import Callable, Optional, Sequence, Union
+
+from PIL import Image
+
+Part = Union[str, Image.Image]
+
+
+@dataclass
+class GenRequest:
+    """Generate text for one user turn of interleaved text/images.
+
+    ``style`` selects the prompt format: ``"chat"`` for the synthesis
+    pipeline's instruction prompts, ``"answer"`` for answering benchmark-style
+    questions (zero-shot and in-context). Backends may format them alike."""
+
+    parts: Sequence[Part]
+    style: str = "chat"
+    max_new_tokens: int = 256
+    temperature: float = 0.0
+    top_p: float = 1.0
+    repetition_penalty: Optional[float] = None   # None -> backend default for the style
+    think: bool = False
+    logprobs: bool = False
+    #: Optional ``fn(partial_text)`` called (from a backend thread) as tokens
+    #: are decoded; backends without streaming simply never call it.
+    on_text: Optional[Callable[[str], None]] = None
+
+
+@dataclass
+class GenResult:
+    text: str                                    # answer, thinking removed
+    raw_text: str = ""
+    token_logprobs: Optional[list[float]] = None
+    num_tokens: int = 0
+
+
+@dataclass
+class ScoreRequest:
+    """Log-probabilities of ``target`` as the answer to ``parts`` (teacher forcing)."""
+
+    parts: Sequence[Part]
+    target: str
+    style: str = "answer"
+
+
+@dataclass
+class ImageRequest:
+    prompt: str
+    width: int = 400
+    height: int = 400
+    seed: Optional[int] = None
+
+
+class Backend:
+    """Base class. Subclasses implement the three ``submit_*`` methods; each
+    returns a ``concurrent.futures.Future`` and may batch freely."""
+
+    name: str = "backend"
+    #: True for unified models that generate images natively.
+    native_image_generation: bool = False
+    #: Default for running the verification critic in thinking mode.
+    verify_think: bool = False
+    #: Defaults for ``SIMITConfig`` fields left as ``None`` (synthesis, use_skills, verify).
+    pipeline_defaults: dict = {}
+
+    def submit_generate(self, req: GenRequest) -> Future:
+        raise NotImplementedError
+
+    def submit_score(self, req: ScoreRequest) -> Future:
+        raise NotImplementedError
+
+    def submit_image(self, req: ImageRequest) -> Future:
+        raise NotImplementedError(f"{self.name} cannot generate images")
+
+    def close(self):
+        pass
+
+    # Synchronous conveniences -------------------------------------------
+    def generate(self, reqs: list[GenRequest]) -> list[GenResult]:
+        futs = [self.submit_generate(r) for r in reqs]
+        return [f.result() for f in futs]
+
+    def score(self, reqs: list[ScoreRequest]) -> list[list[float]]:
+        futs = [self.submit_score(r) for r in reqs]
+        return [f.result() for f in futs]
+
+    def generate_images(self, reqs: list[ImageRequest]) -> list[Image.Image]:
+        futs = [self.submit_image(r) for r in reqs]
+        return [f.result() for f in futs]
+
+
+def chained(fut: Future, fn) -> Future:
+    """A future resolving to ``fn(fut.result())``."""
+    out: Future = Future()
+
+    def done(f):
+        if out.done():
+            return
+        try:
+            out.set_result(fn(f.result()))
+        except BaseException as e:  # propagate engine/postprocess errors to the caller
+            if not out.done():
+                out.set_exception(e)
+
+    fut.add_done_callback(done)
+    out.add_done_callback(lambda o: fut.cancel() if o.cancelled() else None)  # cancellation reaches the engine
+    return out