ollada 0.3.7__py3-none-win_amd64.whl

ollada/__init__.py

@@ -0,0 +1,3 @@
+"""ollada — the shelf for open-source diffusion language models."""
+
+__version__ = "0.3.7"

ollada/adapters/__init__.py

@@ -0,0 +1,33 @@
+"""Family adapter registry."""
+
+from __future__ import annotations
+
+from ollada.adapters.base import AdapterError, BaseAdapter, GenResult
+from ollada.recipes import Recipe
+
+__all__ = ["AdapterError", "BaseAdapter", "GenResult", "create_adapter"]
+
+
+def create_adapter(recipe: Recipe) -> BaseAdapter:
+    """Import lazily so the CLI works on machines without torch installed."""
+    if recipe.family == "dream":
+        from ollada.adapters.dream import DreamAdapter
+
+        return DreamAdapter(recipe)
+    if recipe.family == "sdar":
+        from ollada.adapters.sdar import SdarAdapter
+
+        return SdarAdapter(recipe)
+    if recipe.family == "nemotron":
+        # Dormant: Nemotron is the first planned post-v0.1 addition. Its
+        # adapter stays in the tree; it joins the shelf (catalog) later,
+        # together with the isolated-worker build its transformers>=5
+        # requirement forces (PLAN.md section 5, decision record).
+        from ollada.adapters.nemotron import NemotronAdapter
+
+        return NemotronAdapter(recipe)
+    if recipe.family == "diffcoder":
+        from ollada.adapters.diffcoder import DiffCoderAdapter
+
+        return DiffCoderAdapter(recipe)
+    raise AdapterError(f"no adapter for family '{recipe.family}' (model {recipe.name})")

ollada/adapters/base.py

@@ -0,0 +1,130 @@
+"""Adapter interface — one adapter per model FAMILY.
+
+An adapter hides everything family-specific (how the model loads, what its
+generate() is called, which knobs it takes) behind one interface, so the
+daemon and CLI never see a family's custom code paths directly.
+
+IMPORTANT: an adapter existing is NOT support. Support means the recipe's
+verification block is filled by a real-load run on real GPU hardware
+(PLAN.md section 7). Until then each adapter is unverified.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+
+from ollada.recipes import Recipe
+
+
+@dataclass
+class GenResult:
+    text: str
+    tokens: int
+    seconds: float
+    nfe: int | None = None  # number of function evaluations, when the family reports it
+    # Worker-reported measurement (worker path only; in-process loads use
+    # measurements.capture_and_record instead). None means "not measured".
+    peak_gb: float | None = None
+    hardware: str | None = None
+
+
+class AdapterError(RuntimeError):
+    pass
+
+
+class BaseAdapter:
+    family = "base"
+    # Watch mode: an adapter that can report per-step canvas state
+    # during generation sets this True and accepts on_step in chat().
+    # The manager forwards a step callback only to supporting
+    # adapters; everyone else is called exactly as before.
+    supports_watch = False
+
+    def __init__(self, recipe: Recipe):
+        self.recipe = recipe
+        self.model = None
+        self.tokenizer = None
+
+    # -- lifecycle ---------------------------------------------------------
+    def load(self) -> None:
+        raise NotImplementedError
+
+    def unload(self) -> None:
+        self.model = None
+        self.tokenizer = None
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except Exception:
+            pass
+
+    @property
+    def loaded(self) -> bool:
+        return self.model is not None
+
+    # -- generation --------------------------------------------------------
+    def chat(
+        self, messages: list[dict], overrides: dict | None = None, on_step=None
+    ) -> GenResult:
+        raise NotImplementedError
+
+    # -- helpers ------------------------------------------------------------
+    def _params(self, overrides: dict | None) -> dict:
+        params = dict(self.recipe.generate)
+        if overrides:
+            params.update(overrides)
+        return params
+
+    def _prompt_ids(self, messages: list[dict]):
+        """Tokenize a chat history, preferring the tokenizer's chat template."""
+        from collections.abc import Mapping
+
+        tok = self.tokenizer
+        if getattr(tok, "chat_template", None):
+            result = tok.apply_chat_template(
+                messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
+            )
+            # BatchEncoding is a dict subclass on transformers 4.x but a
+            # UserDict (Mapping, NOT a dict) on 5.x — an isinstance(dict)
+            # check silently passes it through whole on 5.x, and the model's
+            # generate then dies on `.shape` (DiffCoder T4 run, 2026-10-21).
+            if isinstance(result, Mapping):
+                result = result["input_ids"]
+            return result
+        text = ""
+        for msg in messages:
+            text += f"{msg['role']}: {msg['content']}\n"
+        text += "assistant:"
+        return tok(text, return_tensors="pt").input_ids
+
+    def _finish(self, out_ids, prompt_len: int, started: float, nfe: int | None = None) -> GenResult:
+        new_ids = out_ids[0][prompt_len:]
+        text = self.tokenizer.decode(new_ids, skip_special_tokens=True)
+        return GenResult(
+            text=text.strip(),
+            tokens=int(new_ids.shape[0]),
+            seconds=time.time() - started,
+            nfe=nfe,
+        )
+
+    def _require_loaded(self):
+        if not self.loaded:
+            raise AdapterError(f"{self.recipe.name} is not loaded")
+
+
+def quantization_config(quant: str):
+    """Build a BitsAndBytesConfig for int8/int4, else None.
+
+    Never pass load_in_8bit as a raw from_pretrained kwarg — it leaks into
+    some families' custom __init__ and crashes (npc-engine, 2026-10-20).
+    """
+    if quant not in ("int8", "int4"):
+        return None
+    from transformers import BitsAndBytesConfig
+
+    if quant == "int8":
+        return BitsAndBytesConfig(load_in_8bit=True)
+    return BitsAndBytesConfig(load_in_4bit=True)

ollada/adapters/diffcoder.py

@@ -0,0 +1,103 @@
+"""Stable-DiffCoder family adapter (ByteDance-Seed).
+
+Real signature ( modeling_stable_diffcoder.py, read 2026-10-21 ):
+  generate(input_ids=None, generation_config=None, **kwargs) -> plain tensor
+  kwargs forward to generate_block(steps=128, gen_length=128, block_length=4,
+      temperature=0.0, remasking='low_confidence', tokenizer=None,
+      mask_id=5, threshold=0.95, shift=False, eos_id=None)
+The model card's canonical example passes tokenizer=tokenizer,
+threshold=None and eos_id=tokenizer.eos_token_id explicitly — the
+generate_block defaults differ (threshold 0.95, no EOS truncation), so we
+pass all three. Assertions in generate_block: gen_length % block_length == 0
+and steps % (gen_length // block_length) == 0 (recipe defaults satisfy both:
+256 % 4 == 0, 128 % 64 == 0). Batch size must be 1.
+
+The model's config was saved with transformers 5.3.0 — see the recipe pin.
+
+Load repairs applied by this adapter:
+1. generate_block builds its block-causal mask with dtype=torch.bfloat16
+   hardcoded, while the recipe loads fp16 on pre-Ampere GPUs (T4) —
+   transformers 5.x SDPA rejects a float mask whose dtype differs from the
+   query dtype (RuntimeError, T4 run 2026-10-21). The adapter wraps
+   _make_block_causal_mask so the mask follows the model's dtype. The mask
+   is additive (0 / -inf), so fp16 loses nothing.
+
+UNVERIFIED in ollada until a real-load gate passes on a T4.
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ollada.adapters.base import BaseAdapter, quantization_config
+
+_DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16}
+
+
+def _align_mask_dtype(model) -> bool:
+    """Wrap _make_block_causal_mask so its output dtype follows the model.
+
+    Returns True when the wrapper was installed. See the module docstring
+    (repair 1) for why this exists.
+    """
+    orig = getattr(model, "_make_block_causal_mask", None)
+    if orig is None:
+        return False
+    target = model.dtype
+
+    def patched(*args, **kwargs):
+        kwargs["dtype"] = target
+        return orig(*args, **kwargs)
+
+    model._make_block_causal_mask = patched
+    return True
+
+
+class DiffCoderAdapter(BaseAdapter):
+    family = "diffcoder"
+
+    def load(self) -> None:
+        recipe = self.recipe
+        notes: list[str] = []
+        self.tokenizer = AutoTokenizer.from_pretrained(recipe.source, trust_remote_code=True)
+        kwargs = {
+            "trust_remote_code": True,
+            "dtype": _DTYPES.get(recipe.dtype, torch.float16),
+        }
+        qc = quantization_config(recipe.quant_default)
+        if qc is not None:
+            kwargs["quantization_config"] = qc
+            kwargs["device_map"] = "auto"
+        self.model = AutoModelForCausalLM.from_pretrained(recipe.source, **kwargs)
+        if qc is None and torch.cuda.is_available():
+            self.model = self.model.cuda()
+        self.model.eval()
+        if _align_mask_dtype(self.model):
+            notes.append(f"block mask dtype aligned to model dtype ({self.model.dtype})")
+        self.load_notes = notes
+
+    @torch.no_grad()
+    def chat(self, messages, overrides=None):
+        self._require_loaded()
+        params = self._params(overrides)
+        input_ids = self._prompt_ids(messages).to(self.model.device)
+        started = time.time()
+        threshold = params.get("threshold", None)  # card example: None (fixed schedule)
+        out = self.model.generate(
+            input_ids,
+            steps=int(params.get("steps", 128)),
+            gen_length=int(params.get("gen_length", 256)),
+            block_length=int(params.get("block_length", 4)),
+            temperature=float(params.get("temperature", 0.0)),
+            remasking=str(params.get("remasking", "low_confidence")),
+            tokenizer=self.tokenizer,
+            shift=False,
+            threshold=None if threshold is None else float(threshold),
+            eos_id=self.tokenizer.eos_token_id,
+        )
+        # generate() returns a plain tensor (not a GenerateOutput).
+        ids = out.sequences if hasattr(out, "sequences") else out
+        return self._finish(ids, input_ids.shape[1], started)

ollada/adapters/dream.py

@@ -0,0 +1,59 @@
+"""Dream family adapter (Dream-v0-Instruct-7B).
+
+Load path mirrors the proven npc-engine harness: BitsAndBytesConfig for
+8-bit (never a raw load_in_8bit kwarg) and the recipe's transformers 4.x
+pin (5.x breaks Dream's RoPE config with KeyError 'default' — proven live
+on the 5060 Ti run, 2026-10-25; verified on 4.53.3 and 4.57.6).
+
+UNVERIFIED in ollada until the Phase 1 real-load gate passes on a T4.
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+from ollada.adapters.base import BaseAdapter, GenResult, quantization_config
+
+_DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16}
+
+
+class DreamAdapter(BaseAdapter):
+    family = "dream"
+
+    def load(self) -> None:
+        recipe = self.recipe
+        self.tokenizer = AutoTokenizer.from_pretrained(recipe.source, trust_remote_code=True)
+        kwargs = {
+            "trust_remote_code": True,
+            "torch_dtype": _DTYPES.get(recipe.dtype, torch.float16),
+        }
+        qc = quantization_config(recipe.quant_default)
+        if qc is not None:
+            kwargs["quantization_config"] = qc
+            kwargs["device_map"] = "auto"
+        self.model = AutoModel.from_pretrained(recipe.source, **kwargs)
+        if qc is None and torch.cuda.is_available():
+            self.model = self.model.cuda()
+        self.model.eval()
+
+    @torch.no_grad()
+    def chat(self, messages, overrides=None):
+        self._require_loaded()
+        params = self._params(overrides)
+        input_ids = self._prompt_ids(messages).to(self.model.device)
+        started = time.time()
+        out = self.model.diffusion_generate(
+            input_ids,
+            max_new_tokens=int(params.get("gen_length", 256)),
+            steps=int(params.get("steps", 32)),
+            temperature=float(params.get("temperature", 0.2)),
+            top_p=float(params.get("top_p", 0.95)),
+            alg="origin",
+            alg_temp=0.0,
+            mask_id=self.model.config.mask_token_id,
+        )
+        ids = out.sequences if hasattr(out, "sequences") else out
+        return self._finish(ids, input_ids.shape[1], started)

ollada/adapters/nemotron.py

@@ -0,0 +1,70 @@
+"""Nemotron-Labs Diffusion family adapter (NVIDIA).
+
+One checkpoint, three modes — the recipe picks the default:
+  diffusion   -> model.generate(ids, block_length, threshold, eos_token_id)
+  ar          -> model.ar_generate(ids)
+  linear_spec -> model.linear_spec_generate(ids, block_length, eos_token_id)
+Each returns (out_ids, nfe). Requires transformers>=5.0.0 (recipe pin) and
+therefore runs in the isolated nemotron worker env (PLAN.md section 3).
+
+UNVERIFIED in ollada until a real-load gate passes on a T4.
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+from ollada.adapters.base import BaseAdapter, quantization_config
+
+_DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16}
+
+
+class NemotronAdapter(BaseAdapter):
+    family = "nemotron"
+
+    def load(self) -> None:
+        recipe = self.recipe
+        self.tokenizer = AutoTokenizer.from_pretrained(recipe.source, trust_remote_code=True)
+        kwargs = {"trust_remote_code": True}
+        qc = quantization_config(recipe.quant_default)
+        if qc is not None:
+            kwargs["quantization_config"] = qc
+            kwargs["device_map"] = "auto"
+        model = AutoModel.from_pretrained(recipe.source, **kwargs)
+        if qc is None:
+            dtype = _DTYPES.get(recipe.dtype, torch.bfloat16)
+            model = model.to(dtype)
+            if torch.cuda.is_available():
+                model = model.cuda()
+        self.model = model.eval()
+
+    @torch.no_grad()
+    def chat(self, messages, overrides=None):
+        self._require_loaded()
+        params = self._params(overrides)
+        input_ids = self._prompt_ids(messages).to(self.model.device)
+        mode = str(params.get("mode", "diffusion"))
+        max_new = int(params.get("max_new_tokens", 512))
+        eos = self.tokenizer.eos_token_id
+        started = time.time()
+        if mode == "ar":
+            out_ids, nfe = self.model.ar_generate(input_ids, max_new_tokens=max_new)
+        elif mode == "linear_spec":
+            out_ids, nfe = self.model.linear_spec_generate(
+                input_ids,
+                max_new_tokens=max_new,
+                block_length=int(params.get("block_length", 32)),
+                eos_token_id=eos,
+            )
+        else:
+            out_ids, nfe = self.model.generate(
+                input_ids,
+                max_new_tokens=max_new,
+                block_length=int(params.get("block_length", 32)),
+                threshold=float(params.get("threshold", 0.9)),
+                eos_token_id=eos,
+            )
+        return self._finish(out_ids, input_ids.shape[1], started, nfe=int(nfe))