simit 0.1.0__py3-none-any.whl

simit/tune.py

@@ -0,0 +1,323 @@
+"""Hyperparameter search for adaptive budget allocation + difficulty filtering.
+
+Following Appendix G, synthesis runs once: for every validation query we
+cache its zero-shot answer and confidence p0 plus K_max verified candidates
+with their confidences. A trial then only re-selects demonstrations from the
+cache (the first K*(p0) candidates inside the confidence band) and answers.
+Answers are memoized per (query, selected subset) -- at most 2^K_max per
+query -- so later trials are nearly free.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import io
+import json
+import logging
+import pickle
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Optional, Sequence, Union
+
+from PIL import Image
+
+from .config import SIMITConfig
+from .metrics import get_metric
+from .types import Demo, ZeroShot
+from .utils import image_hash, text_hash, to_pil
+
+log = logging.getLogger("simit")
+
+
+@dataclass
+class TuneResult:
+    config: SIMITConfig
+    best_score: float
+    zero_shot_score: float
+    best_params: dict
+    n_trials: int
+    relative_gain: float = 0.0
+    mean_candidates: float = 0.0     # candidates synthesized per query by the best config
+    mean_demos: float = 0.0          # demonstrations shown per query by the best config
+    runtime_ratio: Optional[float] = None   # estimated batched runtime / zero-shot (with max_runtime_ratio)
+    history: list = field(default_factory=list)
+
+    def __repr__(self):
+        rt = f", runtime≈{self.runtime_ratio:.1f}x" if self.runtime_ratio is not None else ""
+        return (f"TuneResult(best={self.best_score:.4f}, zero_shot={self.zero_shot_score:.4f}, "
+                f"gain={100 * self.relative_gain:+.2f}%, candidates/query={self.mean_candidates:.2f}, "
+                f"demos/query={self.mean_demos:.2f}{rt}, params={self.best_params})")
+
+
+def _items(val_set) -> tuple[list[tuple[Image.Image, str, list[str]]], list]:
+    """Validation examples, plus each one's metric override (None: the global one)."""
+    out, metrics = [], []
+    for ex in val_set:
+        metric = None
+        if isinstance(ex, dict):
+            image = ex.get("image")
+            question = ex.get("question") or ex.get("prompt") or ex.get("text")
+            answers = ex.get("answers", ex.get("answer"))
+            metric = ex.get("metric")
+        else:
+            image, question, answers = ex
+        if answers is None:
+            raise ValueError("every validation example needs an answer")
+        answers = [answers] if isinstance(answers, str) else [str(a) for a in answers]
+        out.append((to_pil(image), str(question), answers))
+        metrics.append(get_metric(metric) if metric is not None else None)
+    return out, metrics
+
+
+@dataclass
+class _Entry:
+    zero_shot: ZeroShot
+    candidates: list   # Demo objects, in production order
+
+
+def _cache_key(image, question, sim_id) -> str:
+    return text_hash(image_hash(image), question, sim_id)
+
+
+def _dump(entry: _Entry) -> bytes:
+    cands = []
+    for d in entry.candidates:
+        buf = io.BytesIO()
+        d.image.save(buf, format="PNG")
+        cands.append({**d.to_dict(), "image": buf.getvalue()})
+    return pickle.dumps({"zero_shot": entry.zero_shot, "candidates": cands})
+
+
+def _load(data: bytes) -> _Entry:
+    obj = pickle.loads(data)
+    cands = []
+    for c in obj["candidates"]:
+        img = Image.open(io.BytesIO(c.pop("image")))
+        img.load()
+        cands.append(Demo(image=img, **c))
+    return _Entry(obj["zero_shot"], cands)
+
+
+def select(cfg: SIMITConfig, entry: _Entry) -> tuple[int, ...]:
+    """Indices of the demonstrations SIMIT would use (ABA + DF online rule)."""
+    k = cfg.budget(entry.zero_shot.confidence)
+    if k <= 0:
+        return ()
+    kept = [i for i, d in enumerate(entry.candidates[:cfg.k_max]) if cfg.in_band(d.confidence)]
+    return tuple(kept[:k])
+
+
+def synthesis_cost(cfg: SIMITConfig, entry: _Entry) -> int:
+    """Candidates the online rule synthesizes for this query: it stops once
+    K*(p0) candidates passed DF, or after K_max candidates."""
+    k = cfg.budget(entry.zero_shot.confidence)
+    if k <= 0:
+        return 0
+    cands = entry.candidates[:cfg.k_max]
+    passed = 0
+    for i, d in enumerate(cands):
+        passed += cfg.in_band(d.confidence)
+        if passed >= k:
+            return i + 1
+    return cfg.k_max
+
+
+def reachable_subsets(entry: _Entry, k_max_choices: Sequence[int]) -> set:
+    """Every demo subset ``select`` can return for this query under some
+    configuration: the DF band keeps the candidates whose confidence lies in an
+    interval, and ABA takes the first k of them."""
+    out = set()
+    for km in k_max_choices:
+        cands = entry.candidates[:km]
+        confs = sorted({d.confidence for d in cands if d.confidence is not None})
+        for a, lo in enumerate(confs):
+            for hi in confs[a:]:
+                kept = [i for i, d in enumerate(cands) if d.confidence is not None and lo <= d.confidence <= hi]
+                out.update(tuple(kept[:k]) for k in range(1, km + 1))
+    return out
+
+
+def run_tune(sim, 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) -> TuneResult:
+    import optuna
+
+    items, item_metrics = _items(val_set)
+    default_fn = get_metric(metric)
+
+    def score(i: int, prediction: str) -> float:  # per-example metric for mixed validation sets
+        return (item_metrics[i] or default_fn)(prediction, items[i][2])
+    base_cfg: SIMITConfig = sim.config
+    choices = sorted(set(k_max_choices or [base_cfg.k_max]))
+    k_cache = max(choices)
+    fill_cfg = base_cfg.replace(k_max=k_cache)
+    mnt = max_new_tokens or base_cfg.answer_max_new_tokens
+    pipe = sim._pipe
+    cache_dir = Path(cache_dir) if cache_dir else None
+    if cache_dir:
+        cache_dir.mkdir(parents=True, exist_ok=True)
+    sim_id = text_hash(sim.model_id, str(k_cache), str(mnt), str(fill_cfg.verify), str(fill_cfg.diversity_prompt),
+                       str(fill_cfg.use_skills))
+
+    # ---- 0) batched zero-shot time per query: the unit of max_runtime_ratio
+    zs_list, t_zs = None, None
+    if max_runtime_ratio is not None:
+        async def zs_all():
+            return await asyncio.gather(*[pipe.zero_shot(img, q, mnt) for img, q, _ in items])
+        sim._run(pipe.zero_shot(items[0][0], items[0][1], mnt))  # warm-up (graph capture, kernel compiles)
+        t0 = time.perf_counter()
+        zs_list = sim._run(zs_all())
+        t_zs = (time.perf_counter() - t0) / len(items)
+
+    # ---- 1) synthesis cache: zero-shot + K_max candidates per query
+    entries: list[Optional[_Entry]] = [None] * len(items)
+    todo = []
+    for i, (img, q, _) in enumerate(items):
+        f = cache_dir / f"{_cache_key(img, q, sim_id)}.pkl" if cache_dir else None
+        if f is not None and f.exists():
+            entries[i] = _load(f.read_bytes())
+        else:
+            todo.append(i)
+
+    async def build():
+        sem = asyncio.Semaphore(concurrency)
+        old = pipe.config
+        pipe.config = fill_cfg
+        done = 0
+
+        async def one(i):
+            nonlocal done
+            img, q, _ = items[i]
+            async with sem:
+                res = await pipe.imagine(img, q, fill=True, max_new_tokens=mnt,
+                                         zero_shot=zs_list[i] if zs_list else None)
+            entries[i] = _Entry(res.zero_shot, res.candidates)
+            if cache_dir:
+                (cache_dir / f"{_cache_key(img, q, sim_id)}.pkl").write_bytes(_dump(entries[i]))
+            done += 1
+            log.info("tune cache %d/%d (query %d: %d candidates)", done, len(todo), i, len(res.candidates))
+
+        try:
+            await asyncio.gather(*[one(i) for i in todo])
+        finally:
+            pipe.config = old
+
+    # Batched synthesis seconds per candidate, kept next to the cache for later runs.
+    cost_file = cache_dir / "cost.json" if cache_dir else None
+    spent = json.loads(cost_file.read_text()) if cost_file and cost_file.exists() else {"seconds": 0.0, "candidates": 0}
+    if todo:
+        t0 = time.perf_counter()
+        sim._run(build())
+        spent["seconds"] += time.perf_counter() - t0
+        spent["candidates"] += sum(len(entries[i].candidates) for i in todo)
+        if cost_file:
+            cost_file.write_text(json.dumps(spent))
+    sec_per_candidate = spent["seconds"] / spent["candidates"] if spent["candidates"] else None
+    if max_runtime_ratio is not None and sec_per_candidate is None:
+        log.warning("max_runtime_ratio ignored: no synthesis timing (cache built without it)")
+
+    zs_scores = [score(i, e.zero_shot.answer) for i, e in enumerate(entries)]
+    zero_shot_score = sum(zs_scores) / len(items)
+
+    # ---- 2) trials over cached candidates
+    memo: dict[tuple[int, tuple], float] = {}
+    answered = {"seconds": 0.0, "n": 0}
+
+    def answer(pairs):
+        async def answer_all():
+            sem = asyncio.Semaphore(concurrency * 4)
+
+            async def one(i, sel):
+                img, q, _ = items[i]
+                demos = [entries[i].candidates[j] for j in sel]
+                async with sem:
+                    return await pipe.answer(img, q, demos, mnt)
+            return await asyncio.gather(*[one(i, sel) for i, sel in pairs])
+
+        t0 = time.perf_counter()
+        answers = sim._run(answer_all())
+        answered["seconds"] += time.perf_counter() - t0
+        answered["n"] += len(pairs)
+        for (i, sel), ans in zip(pairs, answers):
+            memo[(i, sel)] = score(i, ans)
+
+    # Only a few subsets per query are reachable: answer them all in one
+    # batch, so trials become lookups (unless that set is unusually large).
+    reach = [(i, sel) for i, e in enumerate(entries) for sel in sorted(reachable_subsets(e, choices))]
+    if len(reach) <= 32 * len(entries):
+        answer(reach)
+
+    def evaluate(cfg: SIMITConfig) -> float:
+        keys = [(i, select(cfg, e)) for i, e in enumerate(entries)]
+        missing = list(dict.fromkeys(k for k in keys if k[1] and k not in memo))
+        if missing:
+            answer(missing)
+        total = 0.0
+        for i, sel in keys:
+            total += memo[(i, sel)] if sel else zs_scores[i]
+        return total / len(items)
+
+    def sample(trial) -> SIMITConfig:
+        eps = trial.suggest_float("epsilon", 0.0, 0.5)
+        a0 = trial.suggest_float("A0", 0.05, 0.6)
+        a1 = trial.suggest_float("A1", a0, 0.9)
+        b = trial.suggest_float("B", 0.0, 0.5)
+        tl = trial.suggest_float("t_low", 0.0, 0.5)
+        th = trial.suggest_float("t_high", tl, 1.0)
+        km = trial.suggest_categorical("k_max", choices) if len(choices) > 1 else choices[0]
+        return base_cfg.replace(epsilon=eps, A0=a0, A1=a1, B=b, t_low=tl, t_high=th, k_max=km,
+                                use_aba=True, use_df=True)
+
+    optuna.logging.set_verbosity(optuna.logging.WARNING)
+    study = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(seed=seed))
+    start = {k: getattr(base_cfg, k) for k in ("epsilon", "A0", "A1", "B", "t_low", "t_high")}
+    if len(choices) > 1:
+        start["k_max"] = base_cfg.k_max if base_cfg.k_max in choices else choices[-1]
+    study.enqueue_trial(start)
+    study.enqueue_trial({**start, "epsilon": 0.0, "t_low": 0.0, "t_high": 1.0})  # plain ICL with K_max demos
+    history = []
+
+    def cost(cfg) -> float:
+        return sum(synthesis_cost(cfg, e) for e in entries) / len(entries)
+
+    def runtime_ratio(cfg) -> Optional[float]:
+        """Zero-shot + synthesis + ICL answers, relative to zero-shot (batched)."""
+        if t_zs is None or sec_per_candidate is None:
+            return None
+        with_demos = sum(bool(select(cfg, e)) for e in entries) / len(entries)
+        per_answer = answered["seconds"] / answered["n"] if answered["n"] else t_zs
+        return 1 + (sec_per_candidate * cost(cfg) + per_answer * with_demos) / t_zs
+
+    def objective(trial):
+        cfg = sample(trial)
+        c = cost(cfg)
+        r = runtime_ratio(cfg) if max_runtime_ratio is not None else None
+        trial.set_user_attr("mean_candidates", c)
+        over = 0.0  # infeasible: over the synthesis / runtime budget
+        if max_candidates is not None and c > max_candidates:
+            over += c - max_candidates
+        if r is not None and r > max_runtime_ratio:
+            over += (r - max_runtime_ratio) / max_runtime_ratio
+        if over > 0:
+            value = -1.0 - over
+        else:
+            value = evaluate(cfg)
+        history.append((trial.params, value, c))
+        return value
+
+    study.optimize(objective, n_trials=n_trials)
+    best_cfg = sample(optuna.trial.FixedTrial(study.best_params))
+    best_score, best_params = study.best_value, study.best_params
+    if best_score <= zero_shot_score:  # nothing (within budget) beats zero-shot: never imagine
+        best_cfg = base_cfg.replace(epsilon=1.0, use_aba=True)
+        best_score, best_params = zero_shot_score, {"epsilon": 1.0}
+    rel = (best_score - zero_shot_score) / zero_shot_score if zero_shot_score > 0 else 0.0
+    if apply:
+        sim.config = best_cfg
+    demos = sum(len(select(best_cfg, e)) for e in entries) / len(entries)
+    return TuneResult(config=best_cfg, best_score=best_score, zero_shot_score=zero_shot_score,
+                      best_params=best_params, n_trials=n_trials, relative_gain=rel,
+                      mean_candidates=cost(best_cfg), mean_demos=demos, runtime_ratio=runtime_ratio(best_cfg),
+                      history=history)

simit/types.py

@@ -0,0 +1,96 @@
+"""Data types returned by SIMIT."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional, Sequence
+
+from PIL import Image
+
+
+@dataclass
+class Demo:
+    """One imagined (image, question, answer) demonstration."""
+
+    image: Image.Image
+    question: str
+    answer: str
+    description: str = ""               # the textual image description it was realized from
+    skill: str = ""                     # skill that realized the image ("natural", "figure", ...)
+    confidence: Optional[float] = None  # model's teacher-forced confidence in ``answer``
+    verify_score: Optional[int] = None  # critic score of the image
+    spec: str = ""                      # skill specification the image was rendered from
+
+    def to_dict(self, image_path: Optional[str] = None) -> dict:
+        return {"image": image_path, "question": self.question, "answer": self.answer,
+                "description": self.description, "skill": self.skill, "confidence": self.confidence,
+                "verify_score": self.verify_score, "spec": self.spec}
+
+    def _repr_html_(self):  # notebooks
+        import base64
+        import io
+        buf = io.BytesIO()
+        img = self.image.copy()
+        img.thumbnail((256, 256))
+        img.save(buf, format="PNG")
+        b64 = base64.b64encode(buf.getvalue()).decode()
+        return (f'<div><img src="data:image/png;base64,{b64}"/><br><b>Q:</b> {self.question}<br>'
+                f'<b>A:</b> {self.answer} <i>({self.skill})</i></div>')
+
+
+@dataclass
+class ZeroShot:
+    answer: str
+    confidence: Optional[float]         # geometric-mean token probability (p0)
+    token_logprobs: Sequence[float] = ()
+
+
+@dataclass
+class Imagination:
+    """Everything ``imagine`` computed for one query."""
+
+    demos: list[Demo]
+    zero_shot: ZeroShot
+    budget: int                          # K*(q) chosen by adaptive budget allocation
+    candidates: list[Demo] = field(default_factory=list)  # verified candidates before filtering
+    stats: dict = field(default_factory=dict)
+
+
+def save_demos(demos: Sequence[Demo], directory, prefix: str = "demo") -> Path:
+    """Write demos as PNGs plus ``demos.jsonl`` (LLaVA-style ``conversations``
+    included), e.g. to fine-tune on them later (SIMIT-FT)."""
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    path = directory / "demos.jsonl"
+    start = 0  # appending: continue numbering so earlier records keep their images
+    while (directory / f"{prefix}_{start}.png").exists():
+        start += 1
+    lines = []
+    for i, d in enumerate(demos, start):
+        name = f"{prefix}_{i}.png"
+        d.image.save(directory / name)
+        rec = d.to_dict(image_path=name)
+        rec["conversations"] = [{"from": "human", "value": f"<image>\n{d.question}"},
+                                {"from": "gpt", "value": d.answer}]
+        lines.append(json.dumps(rec, ensure_ascii=False))
+    with open(path, "a") as f:
+        f.write("\n".join(lines) + ("\n" if lines else ""))
+    return path
+
+
+def load_demos(directory) -> list[Demo]:
+    directory = Path(directory)
+    demos = []
+    for line in (directory / "demos.jsonl").read_text().splitlines():
+        if not line.strip():
+            continue
+        rec = json.loads(line)
+        img = Image.open(directory / rec["image"])
+        img.load()
+        demos.append(Demo(image=img, question=rec["question"], answer=rec["answer"],
+                          description=rec.get("description", ""), skill=rec.get("skill", ""),
+                          confidence=rec.get("confidence"), verify_score=rec.get("verify_score"),
+                          spec=rec.get("spec", "")))
+    return demos

simit/utils.py

@@ -0,0 +1,67 @@
+"""Small shared helpers: image coercion, hashing, response cleanup."""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import re
+from pathlib import Path
+from typing import Union
+
+from PIL import Image
+
+ImageLike = Union[Image.Image, str, Path, bytes]
+
+# Chat special tokens that some backends leave in decoded text.
+_SPECIAL_TOKEN_RE = re.compile(r"<\|im_(?:start|end)\|>|<\|endoftext\|>")
+# CJK ideographs: small models sometimes drift into Chinese mid-output.
+CJK_RE = re.compile(r"[一-鿿]")
+_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
+
+
+def to_pil(image: ImageLike) -> Image.Image:
+    """Load an image from a PIL image, path, or raw bytes, flattened to RGB
+    (transparent pixels become white, matching BAGEL's ``pil_img2rgb``)."""
+    if isinstance(image, (str, Path)):
+        image = Image.open(image)
+    elif isinstance(image, (bytes, bytearray)):
+        image = Image.open(io.BytesIO(image))
+    if not isinstance(image, Image.Image):
+        raise TypeError(f"expected an image, got {type(image).__name__}")
+    if image.mode == "RGBA" or image.info.get("transparency") is not None:
+        image = image.convert("RGBA")
+        white = Image.new("RGB", image.size, (255, 255, 255))
+        white.paste(image, mask=image.split()[3])
+        return white
+    return image.convert("RGB")
+
+
+def image_hash(image: Image.Image) -> str:
+    h = hashlib.sha1()
+    h.update(f"{image.size}{image.mode}".encode())
+    h.update(image.tobytes())
+    return h.hexdigest()
+
+
+def text_hash(*parts: str) -> str:
+    return hashlib.sha1("\x1f".join(parts).encode()).hexdigest()
+
+
+def clean_response(text: str) -> str:
+    return _SPECIAL_TOKEN_RE.sub("", text)
+
+
+def strip_thinking(text: str) -> str:
+    """Drop a leading ``<think>...</think>`` block (or an unterminated one)."""
+    text = _THINK_RE.sub("", text)
+    if "<think>" in text:  # truncated mid-thought: nothing usable after it
+        text = text.split("<think>")[0]
+    if "</think>" in text:  # template already opened the think block
+        text = text.split("</think>", 1)[1]
+    return text.strip()
+
+
+def png_bytes(image: Image.Image) -> bytes:
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    return buf.getvalue()