[mypy] Add mypy type annotation part 1 (#4006)

This commit is contained in:
SangBin Cho
2024-04-13 06:35:50 +09:00
committed by GitHub
parent d4ec9ffb95
commit 09473ee41c
25 changed files with 171 additions and 72 deletions

View File

@@ -171,10 +171,10 @@ class SequenceData:
return self.prompt_token_ids[-1]
return self.output_token_ids[-1]
def get_prompt_token_ids(self) -> int:
def get_prompt_token_ids(self) -> List[int]:
return self.prompt_token_ids
def get_output_token_ids(self) -> int:
def get_output_token_ids(self) -> List[int]:
return self.output_token_ids
@property
@@ -370,7 +370,7 @@ class SequenceGroupState:
"""Mutable state tied to a specific sequence group"""
# torch.Generator used in seeded sampling
generator: Optional = None
generator: Optional = None # type: ignore
class MultiModalData:
@@ -599,7 +599,7 @@ class SequenceGroupMetadata:
return self.lora_request.lora_int_id if self.lora_request else 0
@property
def token_chunk_size(self) -> int:
def token_chunk_size(self) -> Optional[int]:
"""Return the number of tokens to be processed (chunk size)."""
return self._token_chunk_size