Add __repr__

This commit is contained in:
Woosuk Kwon 2023-02-14 09:34:07 +00:00
parent 2729087efe
commit 3363c27d19
2 changed files with 14 additions and 3 deletions

View File

@ -49,3 +49,8 @@ class PhysicalTokenBlock:
self.block_size = block_size
self.ref_count = 0
def __repr__(self) -> str:
return (f'PhysicalTokenBlock(device={self.device}, '
f'block_number={self.block_number}, '
f'ref_count={self.ref_count})')

View File

@ -2,7 +2,6 @@ import enum
from typing import List, Optional
from cacheflow.block import LogicalTokenBlock
from cacheflow.decoding import DecodingParams
class SequenceStatus(enum.Enum):
@ -58,6 +57,11 @@ class Sequence:
token_ids.extend(block.get_token_ids())
return token_ids
def __repr__(self) -> str:
return (f'Sequence(seq_id={self.seq_id}, '
f'status={self.status.name}, '
f'num_blocks={len(self.logical_token_blocks)})')
class SequenceGroup:
@ -65,11 +69,9 @@ class SequenceGroup:
self,
group_id: int,
seqs: List[Sequence],
decoding_params: DecodingParams,
) -> None:
self.group_id = group_id
self.seqs = seqs
self.decoding_params = decoding_params
def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
if status is None:
@ -82,3 +84,7 @@ class SequenceGroup:
if seq.seq_id == seq_id:
return seq
raise ValueError(f'Sequence {seq_id} not found.')
def __repr__(self) -> str:
return (f'SequenceGroup(group_id={self.group_id}, '
f'num_seqs={len(self.seqs)})')