The model output associated with a pooling sequence group.
| 21 | |
| 22 | |
| 23 | class PoolingSequenceGroupOutput( |
| 24 | msgspec.Struct, |
| 25 | omit_defaults=True, |
| 26 | array_like=True, |
| 27 | ): |
| 28 | """The model output associated with a pooling sequence group.""" |
| 29 | |
| 30 | # Annotated as Any to be compatible with msgspec |
| 31 | # The actual type is in SequenceGroup.pooled_data |
| 32 | data: Any |
| 33 | |
| 34 | def get_data_nbytes(self) -> int: |
| 35 | if isinstance(self.data, paddle.Tensor): |
| 36 | return self.data.numel() * self.data.element_size() |
| 37 | elif hasattr(self.data, "nbytes"): |
| 38 | return self.data.nbytes |
| 39 | else: |
| 40 | return 0 |
| 41 | |
| 42 | def __repr__(self) -> str: |
| 43 | return f"PoolingSequenceGroupOutput(data={self.data}" |
| 44 | |
| 45 | def __eq__(self, other: object) -> bool: |
| 46 | if not isinstance(other, PoolingSequenceGroupOutput): |
| 47 | raise NotImplementedError() |
| 48 | return self.data == other.data |
| 49 | |
| 50 | |
| 51 | class PoolerOutput(msgspec.Struct, omit_defaults=True, array_like=True): |
no outgoing calls