token-channel/ref/dspark-b7/dspark.py

# coding=utf-8
"""DSpark draft model built on the DFlash parallel backbone."""

from typing import Optional

import torch
from torch import nn

from .dflash import DFlashDraftModel
from .dspark_heads import AcceptRatePredictor, build_markov_head, sample_tokens


class DSparkDraftModel(DFlashDraftModel):
    """DFlash backbone plus DSpark Markov and confidence heads."""

    def __init__(self, config) -> None:
        self._normalize_dspark_config(config)
        super().__init__(config)
        self.markov_head = build_markov_head(config)
        self.enable_confidence_head = bool(config.enable_confidence_head)
        self.confidence_head_with_markov = False
        if self.enable_confidence_head:
            self.confidence_head_with_markov = bool(config.confidence_head_with_markov)
        if self.enable_confidence_head and self.confidence_head_with_markov:
            assert self.markov_head is not None

        self.confidence_head = None
        if self.enable_confidence_head:
            input_dim = int(config.hidden_size)
            if self.confidence_head_with_markov:
                input_dim += int(config.markov_rank)
            self.confidence_head = AcceptRatePredictor(input_dim=input_dim)
        self.post_init()

    @staticmethod
    def _normalize_dspark_config(config) -> None:
        dspark_config = getattr(config, "dspark_config", None) or {}
        dflash_config = getattr(config, "dflash_config", None) or {}

        for key in ("target_layer_ids", "mask_token_id"):
            if not hasattr(config, key):
                if key in dspark_config:
                    setattr(config, key, dspark_config[key])
                elif key in dflash_config:
                    setattr(config, key, dflash_config[key])
        if not hasattr(config, "num_anchors"):
            config.num_anchors = int(dspark_config.get("num_anchors", 512))
        if not hasattr(config, "markov_rank"):
            config.markov_rank = int(dspark_config.get("markov_rank", 0))
        if int(config.markov_rank) > 0 and not hasattr(config, "markov_head_type"):
            config.markov_head_type = str(dspark_config.get("markov_head_type", "vanilla"))
        if not hasattr(config, "confidence_head_alpha"):
            config.confidence_head_alpha = float(dspark_config.get("confidence_head_alpha", 0.0))
        if not hasattr(config, "enable_confidence_head"):
            config.enable_confidence_head = float(config.confidence_head_alpha) > 0.0
        if bool(config.enable_confidence_head) and not hasattr(config, "confidence_head_with_markov"):
            config.confidence_head_with_markov = bool(
                dspark_config.get("confidence_head_with_markov", False)
            )

        if not hasattr(config, "dflash_config") or config.dflash_config is None:
            config.dflash_config = {}
        if hasattr(config, "target_layer_ids"):
            config.dflash_config.setdefault("target_layer_ids", config.target_layer_ids)
        if hasattr(config, "mask_token_id"):
            config.dflash_config.setdefault("mask_token_id", config.mask_token_id)

    def predict_confidence_step(
        self,
        hidden_states: torch.Tensor,
        prev_token_ids: Optional[torch.Tensor] = None,
    ) -> Optional[torch.Tensor]:
        if self.confidence_head is None:
            return None
        if self.confidence_head_with_markov:
            assert self.markov_head is not None
            assert prev_token_ids is not None
            prev_embeddings = self.markov_head.get_prev_embeddings(prev_token_ids).to(
                dtype=hidden_states.dtype
            )
            features = torch.cat([hidden_states, prev_embeddings], dim=-1)
            return self.confidence_head(features).float()
        return self.confidence_head(hidden_states).float()

    def sample_draft_tokens(
        self,
        base_logits: torch.Tensor,
        *,
        first_prev_token_ids: torch.Tensor,
        temperature: float = 0.0,
        hidden_states: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch_size, proposal_len = base_logits.shape[:2]
        if proposal_len == 0:
            empty_tokens = torch.empty(
                batch_size,
                0,
                dtype=torch.long,
                device=base_logits.device,
            )
            return empty_tokens, base_logits
        if self.markov_head is None:
            return sample_tokens(base_logits, temperature), base_logits
        return self.markov_head.sample_block_tokens(
            base_logits,
            first_prev_token_ids=first_prev_token_ids,
            hidden_states=hidden_states,
            temperature=temperature,
        )

    def sample_draft_token_step(
        self,
        base_logits: torch.Tensor,
        *,
        prev_token_ids: torch.Tensor,
        temperature: float = 0.0,
        hidden_states: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        assert base_logits.ndim == 2, (
            "sample_draft_token_step expects base_logits shaped [batch, vocab], "
            f"got {tuple(base_logits.shape)}."
        )
        if self.markov_head is None:
            step_logits = base_logits
        else:
            step_logits = self.markov_head.apply_step_logits(
                base_logits,
                token_ids=prev_token_ids,
                hidden_states=hidden_states,
            )
        sampled_token_ids = sample_tokens(
            step_logits.unsqueeze(1),
            temperature=temperature,
        ).squeeze(1)
        return sampled_token_ids, step_logits


__all__ = ["DSparkDraftModel"]