Skip to content

chore(api): add type annotations to LLM class and module-level functions#2241

Open
Koushik-Salammagari wants to merge 1 commit into
Lightning-AI:mainfrom
Koushik-Salammagari:type-hints/api
Open

chore(api): add type annotations to LLM class and module-level functions#2241
Koushik-Salammagari wants to merge 1 commit into
Lightning-AI:mainfrom
Koushik-Salammagari:type-hints/api

Conversation

@Koushik-Salammagari

Copy link
Copy Markdown

What does this PR do?

Adds missing Python 3.10+ style type annotations to litgpt/api.py. api.py had the most unannotated functions of any non-utils.py file in the codebase (11 functions), and also contained several implicit-Optional signatures that mypy flags as errors.

Changes

Location Change
LLM.__init__ preprocessor: "Preprocessor | None" = None
LLM.__init__ prompt_style: PromptStyle | None = None
LLM.__init__ devices: int | list[int] | None = None
LLM.__init__ config: Config | None = None
LLM.__init__ checkpoint_dir: Path | None = None
LLM.__init__ fabric: L.Fabric | None = None
LLM.tokenizer property -> Tokenizer
LLM.state_dict destination: dict | None, prefix: str, keep_vars: bool, -> dict
LLM.load_state_dict state_dict: dict, strict: bool, -> Any
LLM.benchmark num_iterations: int, **kwargs: Any, -> tuple[str, dict]
calculate_number_of_devices devices: int | list[int], -> int
benchmark_dict_to_markdown_table data: dict, -> str

The preprocessor parameter uses a forward reference string ("Preprocessor | None") because Preprocessor is defined later in the same file.

No new imports needed — Any was already imported from typing.

Before submitting

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant