Class SamplingInputs

lmdeploy/pytorch/engine/logits_process.py:243–304

@dataclass
class SamplingInputs:
    temperature: torch.Tensor = None
    bad_words: torch.LongTensor = None
    bad_mask: torch.BoolTensor = None
    stop_words: torch.LongTensor = None
    stop_mask: torch.BoolTensor = None
    repetition_penalty: torch.Tensor = None
    top_k: torch.LongTensor = None
    top_p: torch.Tensor = None
    min_p: torch.Tensor = None
    random_seeds: torch.Tensor = None
    random_offsets: torch.Tensor = None
    max_top_k: int = 1
    min_top_p: float = 1.0
    response_formats: list[str, ...] = ()
    logits_processors: list[list[LogitsProcessor]] = None
    max_num_logprobs: None | int = None
    all_ids: None | torch.Tensor = None
    num_ignore_eos: torch.Tensor = None
    batch_size: int = 0
    session_ctx: None | list[dict[str, Any]] = None
    session_to_cleanup: None | list[int] = None
    # for repetition_penalty and ngram
    generated_ids: torch.Tensor | None = None
    generated_ids_cpu: np.ndarray | None = None

    # n gram
    repetition_ngram_size: torch.Tensor | None = None
    repetition_ngram_threshold: torch.Tensor | None = None
    max_repetition_ngram_size: int = 0

    def to_device(self, device: str, non_blocking: bool = False):
        """To device."""
        out_dict = dict()
        if self.generated_ids is None and self.generated_ids_cpu is not None:
            self.generated_ids = torch.from_numpy(self.generated_ids_cpu.copy())
        for f in fields(self):
            k = f.name
            v = getattr(self, k)
            if isinstance(v, torch.Tensor):
                v = v.to(device, non_blocking=non_blocking)
            out_dict[k] = v

        return SamplingInputs(**out_dict)

    def get_delta(self) -> SamplingInputsDelta:
        """Get delta."""
        delta = SamplingInputsDelta()
        for f in fields(self):
            k = f.name
            v = getattr(self, k)
            if isinstance(v, torch.Tensor):
                setattr(delta, k, v)
        return delta

    def update_delta(self, delta: SamplingInputsDelta):
        """Update from delta."""
        for f in fields(delta):