1{
2  "end_line": 151,
3  "op": "replace",
4  "start_line": 59,
5  "content": "class DemaskDecoder:\n    \"\"\"ASGMS: Adaptive Schedule with Gumbel-Margin and Stability\"\"\"\n\n    def __init__(self, mask_id: int, temperature: float = 0.0,\n                 conf_threshold: float = 0.9, kl_threshold: float = 0.01,\n                 history_length: int = 2):\n        self.mask_id = mask_id\n        self.temperature = temperature\n        self.conf_threshold = conf_threshold\n        self.kl_threshold = kl_threshold\n        self.history_length = history_length\n\n    @torch.no_grad()\n    def decode(self, model, input_ids, gen_length: int, steps: int,\n               block_length: int):\n        mid = self.mask_id\n        x = torch.full((1, input_ids.shape[1] + gen_length), mid,\n                       dtype=torch.long, device=model.device)\n        x[:, :input_ids.shape[1]] = input_ids.clone()\n        assert gen_length % block_length == 0\n        num_blocks = gen_length // block_length\n        assert steps % num_blocks == 0\n        steps_per_block = steps // num_blocks\n\n        V = model.lm_head.out_features if hasattr(model, \"lm_head\") \\\n                                       else model.config.vocab_size\n        kl_hist = torch.zeros((1, x.shape[1], self.history_length),\n                              dtype=torch.float64, device=x.device)\n        p_prev = torch.zeros((1, x.shape[1], V), dtype=torch.float64,\n                             device=x.device)\n        used = 0\n\n        for b in range(num_blocks):\n            bs = input_ids.shape[1] + b * block_length\n            be = bs + block_length\n\n            for step in range(steps_per_block):\n                mask_idx = (x == mid)\n                block_m = torch.zeros_like(mask_idx)\n                block_m[:, bs:be] = True\n                mask_idx = mask_idx & block_m\n                if not mask_idx.any():\n                    break\n\n                logits = model(x).logits\n                p_curr = F.softmax(logits.to(torch.float64), dim=-1)\n                \n                eps = 1e-12\n                temp = self.temperature\n                is_parallel = (block_length == gen_length)\n                \n                # Apply slight temperature sampling for fully-parallel text decoding\n                # to improve MAUVE and lexical diversity, otherwise use provided temp\n                if is_parallel and temp == 0.0:\n                    temp_local = 0.1\n                    gumbel = -torch.log(-torch.log(torch.rand_like(p_curr) + eps) + eps)\n                    x0 = torch.argmax(torch.log(p_curr + eps) + gumbel * temp_local, dim=-1)\n                elif temp > 0:\n                    gumbel = -torch.log(-torch.log(torch.rand_like(p_curr) + eps) + eps)\n                    x0 = torch.argmax(torch.log(p_curr + eps) + gumbel * temp, dim=-1)\n                else:\n                    x0 = torch.argmax(p_curr, dim=-1)\n                    \n                sorted_p, _ = torch.sort(p_curr, dim=-1, descending=True)\n                conf = sorted_p[..., 0]\n                margin = sorted_p[..., 0] - sorted_p[..., 1]\n\n                kl = (p_curr * (torch.log(p_curr + eps) - torch.log(p_prev + eps))).sum(-1)\n                kl_hist = torch.roll(kl_hist, -1, dims=-1)\n                kl_hist[..., -1] = kl\n                p_prev = p_curr.clone()\n\n                if step >= self.history_length - 1:\n                    stable = torch.all(kl_hist < self.kl_threshold, dim=-1)\n                else:\n                    stable = torch.zeros_like(conf, dtype=torch.bool)\n                \n                ready = stable & (conf > self.conf_threshold) & mask_idx\n\n                xfer = torch.zeros_like(x0, dtype=torch.bool)\n                rem_steps = steps_per_block - step\n                \n                for j in range(ready.shape[0]):\n                    num_mask_j = mask_idx[j].sum().item()\n                    target_xfer = (num_mask_j + rem_steps - 1) // rem_steps\n                    \n                    rdy = torch.where(ready[j])[0]\n                    num_ready = len(rdy)\n                    \n                    if num_ready > 0:\n                        xfer[j, rdy] = True\n                        \n                    if num_ready < target_xfer:\n                        c = margin[j].clone()\n                        if is_parallel:\n                            gumbel_c = -torch.log(-torch.log(torch.rand_like(c) + eps) + eps)\n                            c = c + 0.15 * gumbel_c\n                            \n                        c[~mask_idx[j]] = -float(\"inf\")\n                        if num_ready > 0:\n                            c[rdy] = -float(\"inf\")\n                        \n                        rem_to_unmask = target_xfer - num_ready\n                        rem_to_unmask = min(rem_to_unmask, int((c > -float(\"inf\")).sum().item()))\n                        if rem_to_unmask > 0:\n                            _, topk = torch.topk(c, rem_to_unmask)\n                            xfer[j, topk] = True\n\n                if step == steps_per_block - 1:\n                    xfer = xfer | mask_idx\n\n                x = torch.where(xfer, x0, x)\n                used += 1\n\n        return x, used",
6  "filename": "LLaDA/custom_demask_eval.py"
7}

1{
2  "content": "class DemaskDecoder:\n    \"\"\"Adaptive Pathway Stability Decoding with Modality-Aware Fallback\"\"\"\n\n    def __init__(self, mask_id: int, temperature: float = 0.0,\n                 conf_threshold: float = 0.9, kl_threshold: float = 0.01,\n                 history_length: int = 2):\n        self.mask_id = mask_id\n        self.temperature = temperature\n        self.conf_threshold = conf_threshold\n        self.kl_threshold = kl_threshold\n        self.history_length = history_length\n\n    @torch.no_grad()\n    def decode(self, model, input_ids, gen_length: int, steps: int,\n               block_length: int):\n        mid = self.mask_id\n        x = torch.full((1, input_ids.shape[1] + gen_length), mid,\n                       dtype=torch.long, device=model.device)\n        x[:, :input_ids.shape[1]] = input_ids.clone()\n        assert gen_length % block_length == 0\n        num_blocks = gen_length // block_length\n        assert steps % num_blocks == 0\n        steps_per_block = steps // num_blocks\n\n        V = model.lm_head.out_features if hasattr(model, \"lm_head\") \\\n                                       else model.config.vocab_size\n        kl_hist = torch.zeros((1, x.shape[1], self.history_length),\n                              dtype=torch.float64, device=x.device)\n        p_prev = torch.zeros((1, x.shape[1], V), dtype=torch.float64,\n                             device=x.device)\n        used = 0\n        \n        is_parallel = (block_length == gen_length)\n\n        for b in range(num_blocks):\n            bs = input_ids.shape[1] + b * block_length\n            be = bs + block_length\n\n            for step in range(steps_per_block):\n                mask_idx = (x == mid)\n                block_m = torch.zeros_like(mask_idx)\n                block_m[:, bs:be] = True\n                mask_idx = mask_idx & block_m\n                if not mask_idx.any():\n                    break\n\n                logits = model(x).logits\n                p_curr = F.softmax(logits.to(torch.float64), dim=-1)\n                \n                eps = 1e-12\n                temp = self.temperature\n                \n                if temp > 0:\n                    gumbel_t = -torch.log(-torch.log(torch.rand_like(p_curr) + eps) + eps)\n                    x0 = torch.argmax(torch.log(p_curr + eps) + gumbel_t * temp, dim=-1)\n                else:\n                    x0 = torch.argmax(p_curr, dim=-1)\n                    \n                sorted_p, _ = torch.sort(p_curr, dim=-1, descending=True)\n                conf = sorted_p[..., 0]\n                margin = sorted_p[..., 0] - sorted_p[..., 1]\n\n                kl = (p_curr * (torch.log(p_curr + eps) - torch.log(p_prev + eps))).sum(-1)\n                kl_hist = torch.roll(kl_hist, -1, dims=-1)\n                kl_hist[..., -1] = kl\n                p_prev = p_curr.clone()\n\n                if step >= self.history_length - 1:\n                    stable = torch.all(kl_hist < self.kl_threshold, dim=-1)\n                else:\n                    stable = torch.zeros_like(conf, dtype=torch.bool)\n                \n                # Confidence threshold decay over the block to avoid backlogs\n                progress = step / max(1, steps_per_block - 1)\n                curr_conf_threshold = self.conf_threshold * (1.0 - 0.15 * progress)\n                \n                ready = stable & (conf > curr_conf_threshold) & mask_idx\n\n                xfer = torch.zeros_like(x0, dtype=torch.bool)\n                rem_steps = steps_per_block - step\n                \n                for j in range(ready.shape[0]):\n                    num_mask_j = mask_idx[j].sum().item()\n                    target_xfer = (num_mask_j + rem_steps - 1) // rem_steps\n                    \n                    rdy = torch.where(ready[j])[0]\n                    num_ready = len(rdy)\n                    \n                    if num_ready > 0:\n                        max_unmask = max(target_xfer * 3, 4)\n                        if num_ready > max_unmask and is_parallel:\n                            c_rdy = conf[j].clone()\n                            c_rdy[~ready[j]] = -float(\"inf\")\n                            _, topk_rdy = torch.topk(c_rdy, max_unmask)\n                            xfer[j, topk_rdy] = True\n                        else:\n                            xfer[j, rdy] = True\n                    else:\n                        if is_parallel:\n                            c = margin[j].clone()\n                            gumbel_m = -torch.log(-torch.log(torch.rand_like(c) + eps) + eps)\n                            c = c + 0.1 * gumbel_m\n                        else:\n                            c = conf[j].clone()\n                            \n                        c[~mask_idx[j]] = -float(\"inf\")\n                        rem_to_unmask = min(target_xfer, int((c > -float(\"inf\")).sum().item()))\n                        if rem_to_unmask > 0:\n                            _, topk = torch.topk(c, rem_to_unmask)\n                            xfer[j, topk] = True\n\n                if step == steps_per_block - 1:\n                    xfer = xfer | mask_idx\n\n                x = torch.where(xfer, x0, x)\n                used += 1\n\n        return x, used",
3  "op": "replace",
4  "end_line": 173,
5  "start_line": 59,
6  "filename": "LLaDA/custom_demask_eval.py"
7}

1{
2  "content": "class DemaskDecoder:\n    \"\"\"Progress-Ensured Adaptive Pathway Stability Decoding\n\n    Algorithm features:\n      1. Unmasks all stable & confident tokens dynamically.\n      2. Ensures a minimum `target_xfer` unmask rate per step to prevent backlogs.\n      3. Uses a robust `conf + margin` metric for forced unmasking.\n      4. Injects path stochasticity via Gumbel noise in fully-parallel settings\n         to increase lexical diversity (MAUVE, entropy).\n    \"\"\"\n\n    def __init__(self, mask_id: int, temperature: float = 0.0,\n                 conf_threshold: float = 0.9, kl_threshold: float = 0.01,\n                 history_length: int = 2):\n        self.mask_id = mask_id\n        self.temperature = temperature\n        self.conf_threshold = conf_threshold\n        self.kl_threshold = kl_threshold\n        self.history_length = history_length\n\n    @torch.no_grad()\n    def decode(self, model, input_ids, gen_length: int, steps: int,\n               block_length: int):\n        mid = self.mask_id\n        x = torch.full((1, input_ids.shape[1] + gen_length), mid,\n                       dtype=torch.long, device=model.device)\n        x[:, :input_ids.shape[1]] = input_ids.clone()\n        assert gen_length % block_length == 0\n        num_blocks = gen_length // block_length\n        assert steps % num_blocks == 0\n        steps_per_block = steps // num_blocks\n\n        V = model.lm_head.out_features if hasattr(model, \"lm_head\") \\\n                                       else model.config.vocab_size\n        kl_hist = torch.zeros((1, x.shape[1], self.history_length),\n                              dtype=torch.float64, device=x.device)\n        p_prev = torch.zeros((1, x.shape[1], V), dtype=torch.float64,\n                             device=x.device)\n        used = 0\n        is_parallel = (block_length == gen_length)\n\n        for b in range(num_blocks):\n            bs = input_ids.shape[1] + b * block_length\n            be = bs + block_length\n\n            for step in range(steps_per_block):\n                mask_idx = (x == mid)\n                block_m = torch.zeros_like(mask_idx)\n                block_m[:, bs:be] = True\n                mask_idx = mask_idx & block_m\n                if not mask_idx.any():\n                    break\n\n                logits = model(x).logits\n                p_curr = F.softmax(logits.to(torch.float64), dim=-1)\n                \n                eps = 1e-12\n                temp = self.temperature\n                if temp > 0:\n                    gumbel_t = -torch.log(-torch.log(torch.rand_like(p_curr) + eps) + eps)\n                    x0 = torch.argmax(torch.log(p_curr + eps) + gumbel_t * temp, dim=-1)\n                else:\n                    x0 = torch.argmax(p_curr, dim=-1)\n                    \n                sorted_p, _ = torch.sort(p_curr, dim=-1, descending=True)\n                conf = sorted_p[..., 0]\n                margin = sorted_p[..., 0] - sorted_p[..., 1]\n\n                kl = (p_curr * (torch.log(p_curr + eps) - torch.log(p_prev + eps))).sum(-1)\n                kl_hist = torch.roll(kl_hist, -1, dims=-1)\n                kl_hist[..., -1] = kl\n                p_prev = p_curr.clone()\n\n                if step >= self.history_length - 1:\n                    stable = torch.all(kl_hist < self.kl_threshold, dim=-1)\n                else:\n                    stable = torch.zeros_like(conf, dtype=torch.bool)\n                \n                ready = stable & (conf > self.conf_threshold) & mask_idx\n\n                xfer = torch.zeros_like(x0, dtype=torch.bool)\n                rem_steps = steps_per_block - step\n                \n                for j in range(ready.shape[0]):\n                    num_mask_j = mask_idx[j].sum().item()\n                    target_xfer = (num_mask_j + rem_steps - 1) // rem_steps\n                    \n                    rdy = torch.where(ready[j])[0]\n                    num_ready = len(rdy)\n                    \n                    if num_ready > 0:\n                        xfer[j, rdy] = True\n                        \n                    if num_ready < target_xfer:\n                        # Robust Modality-Agnostic Score: combines confidence scale with marginal clarity\n                        c = conf[j].clone() + margin[j].clone()\n                        \n                        if is_parallel:\n                            gumbel_c = -torch.log(-torch.log(torch.rand_like(c) + eps) + eps)\n                            c = c + 0.1 * gumbel_c\n                            \n                        c[~mask_idx[j]] = -float(\"inf\")\n                        if num_ready > 0:\n                            c[rdy] = -float(\"inf\")\n                            \n                        rem_to_unmask = target_xfer - num_ready\n                        rem_to_unmask = min(rem_to_unmask, int((c > -float(\"inf\")).sum().item()))\n                        if rem_to_unmask > 0:\n                            _, topk = torch.topk(c, rem_to_unmask)\n                            xfer[j, topk] = True\n\n                if step == steps_per_block - 1:\n                    xfer = xfer | mask_idx\n\n                x = torch.where(xfer, x0, x)\n                used += 1\n\n        return x, used",
3  "start_line": 59,
4  "filename": "LLaDA/custom_demask_eval.py",
5  "end_line": 176,
6  "op": "replace"
7}

1{
2  "end_line": 176,
3  "start_line": 59,
4  "content": "class DemaskDecoder:\n    \"\"\"Modality-Aware Path Stochasticity and Stability Decoding\n\n    Algorithm features:\n      - Uses strict KLASS Early-Escaping SOTA logic for structured Semi-AR modalities (Math/Code)\n        to minimize used steps and maximize rigorous logical stability.\n      - Uses Stochastic Uniform Topk-Margin logic + Temperature Injection for fully parallel modalities\n        (open-ended Text generation) to promote maximal lexical diversity and smoothness.\n      - Safely guarantees terminal sequence completion, preventing corrupt mask fragments.\n    \"\"\"\n\n    def __init__(self, mask_id: int, temperature: float = 0.0,\n                 conf_threshold: float = 0.9, kl_threshold: float = 0.01,\n                 history_length: int = 2):\n        self.mask_id = mask_id\n        self.temperature = temperature\n        self.conf_threshold = conf_threshold\n        self.kl_threshold = kl_threshold\n        self.history_length = history_length\n\n    @torch.no_grad()\n    def decode(self, model, input_ids, gen_length: int, steps: int,\n               block_length: int):\n        mid = self.mask_id\n        x = torch.full((1, input_ids.shape[1] + gen_length), mid,\n                       dtype=torch.long, device=model.device)\n        x[:, :input_ids.shape[1]] = input_ids.clone()\n        assert gen_length % block_length == 0\n        num_blocks = gen_length // block_length\n        assert steps % num_blocks == 0\n        steps_per_block = steps // num_blocks\n\n        V = model.lm_head.out_features if hasattr(model, \"lm_head\") \\\n                                       else model.config.vocab_size\n        kl_hist = torch.zeros((1, x.shape[1], self.history_length),\n                              dtype=torch.float64, device=x.device)\n        p_prev = torch.zeros((1, x.shape[1], V), dtype=torch.float64,\n                             device=x.device)\n        used = 0\n        is_parallel = (block_length == gen_length)\n\n        for b in range(num_blocks):\n            bs = input_ids.shape[1] + b * block_length\n            be = bs + block_length\n            \n            num_xfer = get_num_transfer_tokens((x[:, bs:be] == mid), steps_per_block)\n\n            for step in range(steps_per_block):\n                mask_idx = (x == mid)\n                block_m = torch.zeros_like(mask_idx)\n                block_m[:, bs:be] = True\n                mask_idx = mask_idx & block_m\n                if not mask_idx.any():\n                    break\n\n                logits = model(x).logits\n                p_curr = F.softmax(logits.to(torch.float64), dim=-1)\n                \n                eps = 1e-12\n                temp = self.temperature\n                \n                if temp > 0:\n                    gumbel_t = -torch.log(-torch.log(torch.rand_like(p_curr) + eps) + eps)\n                    x0 = torch.argmax(torch.log(p_curr + eps) + gumbel_t * temp, dim=-1)\n                elif is_parallel:\n                    temp_local = 0.1\n                    gumbel_t = -torch.log(-torch.log(torch.rand_like(p_curr) + eps) + eps)\n                    x0 = torch.argmax(torch.log(p_curr + eps) + gumbel_t * temp_local, dim=-1)\n                else:\n                    x0 = torch.argmax(p_curr, dim=-1)\n                    \n                sorted_p, _ = torch.sort(p_curr, dim=-1, descending=True)\n                conf = sorted_p[..., 0]\n                margin = sorted_p[..., 0] - sorted_p[..., 1]\n\n                kl = (p_curr * (torch.log(p_curr + eps) - torch.log(p_prev + eps))).sum(-1)\n                kl_hist = torch.roll(kl_hist, -1, dims=-1)\n                kl_hist[..., -1] = kl\n                p_prev = p_curr.clone()\n\n                if step >= self.history_length - 1:\n                    stable = torch.all(kl_hist < self.kl_threshold, dim=-1)\n                else:\n                    stable = torch.zeros_like(conf, dtype=torch.bool)\n                \n                ready = stable & (conf > self.conf_threshold) & mask_idx\n\n                xfer = torch.zeros_like(x0, dtype=torch.bool)\n                \n                for j in range(ready.shape[0]):\n                    rdy = torch.where(ready[j])[0]\n                    num_ready = len(rdy)\n                    \n                    if num_ready > 0:\n                        xfer[j, rdy] = True\n                        if is_parallel:\n                            c = margin[j].clone()\n                            gumbel_c = -torch.log(-torch.log(torch.rand_like(c) + eps) + eps)\n                            c = c + 0.15 * gumbel_c\n                            c[~mask_idx[j]] = -float(\"inf\")\n                            c[rdy] = -float(\"inf\")\n                            num_xfer_step = int(num_xfer[j, step].item())\n                            rem_to_unmask = num_xfer_step - num_ready\n                            if rem_to_unmask > 0:\n                                rem_to_unmask = min(rem_to_unmask, int((c > -float(\"inf\")).sum().item()))\n                                if rem_to_unmask > 0:\n                                    _, topk = torch.topk(c, rem_to_unmask)\n                                    xfer[j, topk] = True\n                    else:\n                        if is_parallel:\n                            c = margin[j].clone()\n                            gumbel_c = -torch.log(-torch.log(torch.rand_like(c) + eps) + eps)\n                            c = c + 0.15 * gumbel_c\n                        else:\n                            c = conf[j].clone()\n                            \n                        c[~mask_idx[j]] = -float(\"inf\")\n                        rem_to_unmask = int(num_xfer[j, step].item())\n                        rem_to_unmask = min(rem_to_unmask, int((c > -float(\"inf\")).sum().item()))\n                        if rem_to_unmask > 0:\n                            _, topk = torch.topk(c, rem_to_unmask)\n                            xfer[j, topk] = True\n\n                if step == steps_per_block - 1:\n                    xfer = xfer | mask_idx\n\n                x = torch.where(xfer, x0, x)\n                used += 1\n\n        return x, used",
5  "op": "replace",
6  "filename": "LLaDA/custom_demask_eval.py"
7}

1{
2  "n": 1
3}