mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-26 16:35:05 +00:00
22 lines
619 B
Python
22 lines
619 B
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Self
|
|
|
|
import torch
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CausalLMForwardOutput:
|
|
hidden_states: torch.Tensor
|
|
"""Input of the lm_head. Shape: [num_tokens, hidden_size]."""
|
|
|
|
def narrow(self, num_tokens: int) -> Self:
|
|
assert 0 < num_tokens <= self.hidden_states.size(0), f"{num_tokens=}, {self.hidden_states.size(0)=}"
|
|
return type(self)(
|
|
hidden_states=self.hidden_states[:num_tokens],
|
|
)
|
|
|
|
def copy_(self, src: Self) -> None:
|
|
self.hidden_states.copy_(src.hidden_states)
|