Spaces:

Ruurd
/

llada2.2-mini

Sleeping

Deploy LLaDA2.2-mini with native edits and BG + TG visualization

b30553b verified 12 days ago

12 kB

"""Observe the pinned native decoder without replacing its sampling/edit logic."""
from __future__ import annotations

from dataclasses import dataclass
from html import escape
from queue import Empty, Full, Queue
from threading import Event, Lock, Thread
import time

import torch

MODEL_ID = "inclusionAI/LLaDA2.2-mini"
MODEL_REVISION = "e4f77d25ab5ca86f2223c6cdc6a72b887e20c8a5"
MASK, DELETE, INSERT, EOS = 156895, 156930, 156931, 156892


@dataclass(frozen=True)
class Mark:
    probability: float = 0.0
    iteration: int = 0


@dataclass
class Frame:
    tokens: list[int]
    marks: list[Mark]
    iteration: int
    deletes: int = 0
    inserts: int = 0
    done: bool = False


def move_marks(written, old, marks, probabilities, iteration, length):
    """Mirror native DELETE / [MASK, old token] INSERT and fixed-block resizing."""
    result = []
    for i, token in enumerate(written):
        if token == DELETE:
            continue
        if token == INSERT:
            result.extend([Mark(), marks[i]])
        else:
            result.append(marks[i] if token == old[i] else Mark(probabilities[i], iteration))
    return (result + [Mark()] * length)[:length]


class NativeStreamer:
    def __init__(self, model):
        self.model = model
        self.lock = Lock()

    def stream(self, input_ids, *, timeout=240, selector=None, **options):
        """Run unchanged model.generate; send CPU snapshots after native edit operations.
        A bounded queue applies backpressure. Cancellation and errors restore instance methods
        before another request can enter; no patched global classes or retained GPU frames.
        """
        queue = Queue(maxsize=2)
        cancel = Event()

        def send(item):
            while not cancel.is_set():
                try:
                    queue.put(item, timeout=.1)
                    return
                except Full:
                    pass

        def work():
            try:
                with self.lock, torch.inference_mode():
                    if cancel.is_set():
                        return
                    self._run(input_ids, options, send, cancel, timeout, selector)
            except BaseException as exc:
                send(exc)
            finally:
                send(None)

        thread = Thread(target=work, daemon=True)
        thread.start()
        try:
            while True:
                try:
                    item = queue.get(timeout=.1)
                except Empty:
                    if not thread.is_alive():
                        break
                    continue
                if item is None:
                    break
                if isinstance(item, BaseException):
                    raise item
                yield item
        finally:
            cancel.set()
            thread.join()

    def _run(self, inputs, options, send, cancel, timeout, selector):
        model = self.model
        start = time.monotonic()
        prompt_length = inputs.shape[1]
        state = dict(iteration=0, deletes=0, inserts=0, marks=[], logits=None)
        # Save instance attributes separately: restoring a bound method onto an instance
        # would shadow its original class descriptor indefinitely.
        names = ("_joint_decode_block", "_apply_edit_operations_with_tracking", "_diffusion_sample")
        saved = {name: model.__dict__.get(name) for name in names}
        present = {name: name in model.__dict__ for name in names}
        joint = model._joint_decode_block
        edit = model._apply_edit_operations_with_tracking
        sample = model._diffusion_sample

        def guided_sample(logits, temperature=1.0, top_k=None, top_p=None):
            tokens, probabilities = sample(logits, temperature, top_k, top_p)
            state["sample_calls"] += 1
            initial_calls = 2 if options["temperature"] != 0 else 1
            if selector is None or state["sample_calls"] > initial_calls:
                return tokens, probabilities
            old = state["x"][0, state["block_start"]:state["block_end"]]
            # Edit decisions remain native. Guide only contiguous ordinary-token
            # predictions, without scoring any edge across a structural operation.
            allowed = (tokens != DELETE) & (tokens != INSERT) & (tokens != MASK)
            positions = torch.arange(len(tokens), device=tokens.device) + state["block_start"]
            allowed &= positions >= prompt_length
            begin = 0
            while begin < len(tokens):
                if not allowed[begin]:
                    begin += 1
                    continue
                end = begin + 1
                while end < len(tokens) and allowed[end]:
                    end += 1
                scores = logits[begin:end].clone()
                scores[:, [MASK, DELETE, INSERT]] = -torch.inf
                absolute = state["block_start"] + begin
                left_id = None
                # Only completed blocks or immutable prompt tokens supply a
                # boundary. Never bridge an edit barrier inside the active block.
                if absolute > 0 and (begin == 0 or absolute == prompt_length):
                    previous = state["x"][0, absolute - 1].item()
                    if previous not in (MASK, DELETE, INSERT):
                        left_id = previous
                selected, _ = selector.select_block(
                    scores, old[begin:end], torch.ones(end-begin, dtype=torch.bool, device=tokens.device),
                    top_k=top_k or 16, temperature=temperature, left_id=left_id,
                )
                tokens[begin:end] = selected
                begin = end
            # Confidence gates always use the full, unmodified model distribution.
            probabilities = logits.float().softmax(-1).gather(-1, tokens[:, None]).squeeze(-1)
            return tokens, probabilities

        def before_forward(module, args, kwargs):
            if cancel.is_set():
                raise RuntimeError("Generation cancelled")
            if time.monotonic() - start > timeout:
                raise RuntimeError("GPU time budget reached; generation is incomplete. Try fewer tokens or refinement steps.")

        def after_forward(module, args, output):
            state["logits"] = output.logits
            state["sample_calls"] = 0

        def observe_block(x, block_start, block_end, *args, **kwargs):
            state.update(x=x, block_start=block_start, block_end=block_end)
            if not state["marks"]:
                state["marks"] = [Mark()] * x.shape[1]
            return joint(x, block_start, block_end, *args, **kwargs)

        def observe_edit(written, old, tracking, mask_id, block_length, delete_token_id, split_token_id):
            result, new_tracking = edit(written, old, tracking, mask_id, block_length,
                                        delete_token_id, split_token_id)
            state["iteration"] += 1
            lo, hi = state["block_start"], state["block_end"]
            logits = state["logits"][0, lo:hi].float()
            ids = torch.tensor(written, device=logits.device).unsqueeze(-1)
            probabilities = logits.softmax(-1).gather(-1, ids).squeeze(-1).tolist()
            state["logits"] = None
            state["marks"][lo:hi] = move_marks(written, old, state["marks"][lo:hi],
                                               probabilities, state["iteration"], block_length)
            state["deletes"] += written.count(DELETE)
            state["inserts"] += written.count(INSERT)
            tokens = state["x"][0].tolist()
            tokens[lo:hi] = result
            end = prompt_length + options["gen_length"]
            send(Frame(tokens[prompt_length:end], list(state["marks"][prompt_length:end]),
                       state["iteration"], state["deletes"], state["inserts"]))
            return result, new_tracking

        try:
            model._joint_decode_block = observe_block
            model._apply_edit_operations_with_tracking = observe_edit
            model._diffusion_sample = guided_sample
            # Native decoder calls self.forward directly, bypassing Module.__call__ hooks.
            original_forward = model.forward
            forward_present = "forward" in model.__dict__
            saved_forward = model.__dict__.get("forward")
            def observed_forward(*args, **kwargs):
                before_forward(model, args, kwargs)
                output = original_forward(*args, **kwargs)
                after_forward(model, args, output)
                return output
            model.forward = observed_forward
            try:
                output = model.generate(inputs=inputs, **options)
            finally:
                if forward_present:
                    model.forward = saved_forward
                else:
                    del model.forward
            tokens = output[0].tolist()
            if any(token in (MASK, DELETE, INSERT) for token in tokens):
                raise RuntimeError("Native decoder reached its iteration cap with unresolved tokens. Increase the cap or reduce output length.")
            send(Frame(tokens, list(state["marks"][prompt_length:prompt_length + len(tokens)]),
                       state["iteration"], state["deletes"], state["inserts"], done=True))
        finally:
            for name in names:
                if present[name]:
                    setattr(model, name, saved[name])
                else:
                    delattr(model, name)


def render(frame, tokenizer, color_mode="Prediction iteration"):
    tokens = frame.tokens
    if EOS in tokens:
        tokens = tokens[:tokens.index(EOS)]
    pieces = []
    for i, token in enumerate(tokens):
        mark = frame.marks[i] if i < len(frame.marks) else Mark()
        text = "mask" if token == MASK else tokenizer.decode([token], skip_special_tokens=False)
        if token == MASK:
            style = "background:#d1d5db;color:#4b5563;border-radius:4px;padding:0 4px;margin:0 1px"
        elif color_mode == "Prediction probability":
            style = f"color:hsl({int(mark.probability * 120)},90%,30%)"
        elif color_mode == "Prediction iteration":
            lightness = 72 - round(42 * mark.iteration / max(frame.iteration, 1))
            style = f"color:hsl(210,90%,{lightness}%)"
        else:
            style = "color:inherit"
        title = f"position {i} · predicted at iteration {mark.iteration} · model probability {mark.probability:.1%}"
        pieces.append(f'<span style="{style}" title="{title}">{escape(text)}</span>')
    # Decode contiguous spans, preserving whitespace and multi-byte token sequences.
    spans, pending = [], []
    for token in tokens + [MASK]:
        if token == MASK:
            spans.append(tokenizer.decode(pending, skip_special_tokens=True, clean_up_tokenization_spaces=False))
            spans.append("[MASK]")
            pending = []
        else:
            pending.append(token)
    plain = "".join(spans[:-1])
    return ("<div style='font-family:system-ui;padding:14px;border:1px solid #d1d5db;"
            "border-radius:9px;background:#fafafa'>"
            f"<div style='font-weight:700;color:#2563eb'>Denoising iteration {frame.iteration} · "
            f"{len(tokens)} output positions · {frame.inserts} insertions · {frame.deletes} deletions</div>"
            f"<div style='line-height:2;font-size:15px;white-space:pre-wrap'>{''.join(pieces)}</div>"
            "<div style='font-size:11px;color:#6b7280'>Gray chips are masks. Hover for prediction details. "
            "Inserted and surviving tokens carry their prediction history as positions shift.</div>"
            "<details style='margin-top:10px'><summary>Plain text / copy code</summary>"
            "<textarea readonly aria-label='Plain text output' spellcheck='false' "
            "style='width:100%;min-height:9em;white-space:pre;font-family:monospace'>"
            f"{escape(plain)}</textarea></details></div>")

	"""Observe the pinned native decoder without replacing its sampling/edit logic."""
	from __future__ import annotations

	from dataclasses import dataclass
	from html import escape
	from queue import Empty, Full, Queue
	from threading import Event, Lock, Thread
	import time

	import torch

	MODEL_ID = "inclusionAI/LLaDA2.2-mini"
	MODEL_REVISION = "e4f77d25ab5ca86f2223c6cdc6a72b887e20c8a5"
	MASK, DELETE, INSERT, EOS = 156895, 156930, 156931, 156892


	@dataclass(frozen=True)
	class Mark:
	probability: float = 0.0
	iteration: int = 0


	@dataclass
	class Frame:
	tokens: list[int]
	marks: list[Mark]
	iteration: int
	deletes: int = 0
	inserts: int = 0
	done: bool = False


	def move_marks(written, old, marks, probabilities, iteration, length):
	"""Mirror native DELETE / [MASK, old token] INSERT and fixed-block resizing."""
	result = []
	for i, token in enumerate(written):
	if token == DELETE:
	continue
	if token == INSERT:
	result.extend([Mark(), marks[i]])
	else:
	result.append(marks[i] if token == old[i] else Mark(probabilities[i], iteration))
	return (result + [Mark()] * length)[:length]


	class NativeStreamer:
	def __init__(self, model):
	self.model = model
	self.lock = Lock()

	def stream(self, input_ids, , timeout=240, selector=None, *options):
	"""Run unchanged model.generate; send CPU snapshots after native edit operations.

	A bounded queue applies backpressure. Cancellation and errors restore instance methods
	before another request can enter; no patched global classes or retained GPU frames.
	"""
	queue = Queue(maxsize=2)
	cancel = Event()

	def send(item):
	while not cancel.is_set():
	try:
	queue.put(item, timeout=.1)
	return
	except Full:
	pass

	def work():
	try:
	with self.lock, torch.inference_mode():
	if cancel.is_set():
	return
	self._run(input_ids, options, send, cancel, timeout, selector)
	except BaseException as exc:
	send(exc)
	finally:
	send(None)

	thread = Thread(target=work, daemon=True)
	thread.start()
	try:
	while True:
	try:
	item = queue.get(timeout=.1)
	except Empty:
	if not thread.is_alive():
	break
	continue
	if item is None:
	break
	if isinstance(item, BaseException):
	raise item
	yield item
	finally:
	cancel.set()
	thread.join()

	def _run(self, inputs, options, send, cancel, timeout, selector):
	model = self.model
	start = time.monotonic()
	prompt_length = inputs.shape[1]
	state = dict(iteration=0, deletes=0, inserts=0, marks=[], logits=None)
	# Save instance attributes separately: restoring a bound method onto an instance
	# would shadow its original class descriptor indefinitely.
	names = ("_joint_decode_block", "_apply_edit_operations_with_tracking", "_diffusion_sample")
	saved = {name: model.__dict__.get(name) for name in names}
	present = {name: name in model.__dict__ for name in names}
	joint = model._joint_decode_block
	edit = model._apply_edit_operations_with_tracking
	sample = model._diffusion_sample

	def guided_sample(logits, temperature=1.0, top_k=None, top_p=None):
	tokens, probabilities = sample(logits, temperature, top_k, top_p)
	state["sample_calls"] += 1
	initial_calls = 2 if options["temperature"] != 0 else 1
	if selector is None or state["sample_calls"] > initial_calls:
	return tokens, probabilities
	old = state["x"][0, state["block_start"]:state["block_end"]]
	# Edit decisions remain native. Guide only contiguous ordinary-token
	# predictions, without scoring any edge across a structural operation.
	allowed = (tokens != DELETE) & (tokens != INSERT) & (tokens != MASK)
	positions = torch.arange(len(tokens), device=tokens.device) + state["block_start"]
	allowed &= positions >= prompt_length
	begin = 0
	while begin < len(tokens):
	if not allowed[begin]:
	begin += 1
	continue
	end = begin + 1
	while end < len(tokens) and allowed[end]:
	end += 1
	scores = logits[begin:end].clone()
	scores[:, [MASK, DELETE, INSERT]] = -torch.inf
	absolute = state["block_start"] + begin
	left_id = None
	# Only completed blocks or immutable prompt tokens supply a
	# boundary. Never bridge an edit barrier inside the active block.
	if absolute > 0 and (begin == 0 or absolute == prompt_length):
	previous = state["x"][0, absolute - 1].item()
	if previous not in (MASK, DELETE, INSERT):
	left_id = previous
	selected, _ = selector.select_block(
	scores, old[begin:end], torch.ones(end-begin, dtype=torch.bool, device=tokens.device),
	top_k=top_k or 16, temperature=temperature, left_id=left_id,
	)
	tokens[begin:end] = selected
	begin = end
	# Confidence gates always use the full, unmodified model distribution.
	probabilities = logits.float().softmax(-1).gather(-1, tokens[:, None]).squeeze(-1)
	return tokens, probabilities

	def before_forward(module, args, kwargs):
	if cancel.is_set():
	raise RuntimeError("Generation cancelled")
	if time.monotonic() - start > timeout:
	raise RuntimeError("GPU time budget reached; generation is incomplete. Try fewer tokens or refinement steps.")

	def after_forward(module, args, output):
	state["logits"] = output.logits
	state["sample_calls"] = 0

	def observe_block(x, block_start, block_end, args, *kwargs):
	state.update(x=x, block_start=block_start, block_end=block_end)
	if not state["marks"]:
	state["marks"] = [Mark()] * x.shape[1]
	return joint(x, block_start, block_end, args, *kwargs)

	def observe_edit(written, old, tracking, mask_id, block_length, delete_token_id, split_token_id):
	result, new_tracking = edit(written, old, tracking, mask_id, block_length,
	delete_token_id, split_token_id)
	state["iteration"] += 1
	lo, hi = state["block_start"], state["block_end"]
	logits = state["logits"][0, lo:hi].float()
	ids = torch.tensor(written, device=logits.device).unsqueeze(-1)
	probabilities = logits.softmax(-1).gather(-1, ids).squeeze(-1).tolist()
	state["logits"] = None
	state["marks"][lo:hi] = move_marks(written, old, state["marks"][lo:hi],
	probabilities, state["iteration"], block_length)
	state["deletes"] += written.count(DELETE)
	state["inserts"] += written.count(INSERT)
	tokens = state["x"][0].tolist()
	tokens[lo:hi] = result
	end = prompt_length + options["gen_length"]
	send(Frame(tokens[prompt_length:end], list(state["marks"][prompt_length:end]),
	state["iteration"], state["deletes"], state["inserts"]))
	return result, new_tracking

	try:
	model._joint_decode_block = observe_block
	model._apply_edit_operations_with_tracking = observe_edit
	model._diffusion_sample = guided_sample
	# Native decoder calls self.forward directly, bypassing Module.__call__ hooks.
	original_forward = model.forward
	forward_present = "forward" in model.__dict__
	saved_forward = model.__dict__.get("forward")
	def observed_forward(args, *kwargs):
	before_forward(model, args, kwargs)
	output = original_forward(args, *kwargs)
	after_forward(model, args, output)
	return output
	model.forward = observed_forward
	try:
	output = model.generate(inputs=inputs, **options)
	finally:
	if forward_present:
	model.forward = saved_forward
	else:
	del model.forward
	tokens = output[0].tolist()
	if any(token in (MASK, DELETE, INSERT) for token in tokens):
	raise RuntimeError("Native decoder reached its iteration cap with unresolved tokens. Increase the cap or reduce output length.")
	send(Frame(tokens, list(state["marks"][prompt_length:prompt_length + len(tokens)]),
	state["iteration"], state["deletes"], state["inserts"], done=True))
	finally:
	for name in names:
	if present[name]:
	setattr(model, name, saved[name])
	else:
	delattr(model, name)


	def render(frame, tokenizer, color_mode="Prediction iteration"):
	tokens = frame.tokens
	if EOS in tokens:
	tokens = tokens[:tokens.index(EOS)]
	pieces = []
	for i, token in enumerate(tokens):
	mark = frame.marks[i] if i < len(frame.marks) else Mark()
	text = "mask" if token == MASK else tokenizer.decode([token], skip_special_tokens=False)
	if token == MASK:
	style = "background:#d1d5db;color:#4b5563;border-radius:4px;padding:0 4px;margin:0 1px"
	elif color_mode == "Prediction probability":
	style = f"color:hsl({int(mark.probability * 120)},90%,30%)"
	elif color_mode == "Prediction iteration":
	lightness = 72 - round(42 * mark.iteration / max(frame.iteration, 1))
	style = f"color:hsl(210,90%,{lightness}%)"
	else:
	style = "color:inherit"
	title = f"position {i} · predicted at iteration {mark.iteration} · model probability {mark.probability:.1%}"
	pieces.append(f'<span style="{style}" title="{title}">{escape(text)}</span>')
	# Decode contiguous spans, preserving whitespace and multi-byte token sequences.
	spans, pending = [], []
	for token in tokens + [MASK]:
	if token == MASK:
	spans.append(tokenizer.decode(pending, skip_special_tokens=True, clean_up_tokenization_spaces=False))
	spans.append("[MASK]")
	pending = []
	else:
	pending.append(token)
	plain = "".join(spans[:-1])
	return ("<div style='font-family:system-ui;padding:14px;border:1px solid #d1d5db;"
	"border-radius:9px;background:#fafafa'>"
	f"<div style='font-weight:700;color:#2563eb'>Denoising iteration {frame.iteration} · "
	f"{len(tokens)} output positions · {frame.inserts} insertions · {frame.deletes} deletions</div>"
	f"<div style='line-height:2;font-size:15px;white-space:pre-wrap'>{''.join(pieces)}</div>"
	"<div style='font-size:11px;color:#6b7280'>Gray chips are masks. Hover for prediction details. "
	"Inserted and surviving tokens carry their prediction history as positions shift.</div>"
	"<details style='margin-top:10px'><summary>Plain text / copy code</summary>"
	"<textarea readonly aria-label='Plain text output' spellcheck='false' "
	"style='width:100%;min-height:9em;white-space:pre;font-family:monospace'>"
	f"{escape(plain)}</textarea></details></div>")