better type annotations (#242)

Signed-off-by: Mayank Mishra <mayank31398@gmail.com>

Author: mayank31398

lm_engine/checkpointing/lr_scheduler.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+from __future__ import annotations
+
 import os
 
 import torch.nn as nn
@@ -16,10 +18,10 @@
 
 
 class _LRSchedulerSaver(Stateful):
-    def __init__(self, lr_scheduler_container: LRSchedulerContainer) -> None:
+    def __init__(self, lr_scheduler_container: LRSchedulerContainer) -> _LRSchedulerSaver:
         self.lr_scheduler_container = lr_scheduler_container
 
-    def state_dict(self) -> dict:
+    def state_dict(self) -> list[dict]:
         return [lr_scheduler.state_dict() for lr_scheduler in self.lr_scheduler_container]
 
     def load_state_dict(self, state_dict: list[dict]) -> None:

lm_engine/checkpointing/model.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+from __future__ import annotations
+
 import os
 
 from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict, set_model_state_dict
@@ -11,7 +13,7 @@
 
 
 class _ModelSaver(Stateful):
-    def __init__(self, model_container: ModelContainer) -> None:
+    def __init__(self, model_container: ModelContainer) -> _ModelSaver:
         self.model_container = model_container
 
     def state_dict(self) -> dict:

lm_engine/checkpointing/model_optimizer.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+from __future__ import annotations
+
 import os
 
 from torch.distributed.checkpoint.state_dict import (
@@ -17,7 +19,9 @@
 
 
 class _ModelOptimizerSaver(Stateful):
-    def __init__(self, model_container: ModelContainer, optimizer_container: OptimizerContainer) -> None:
+    def __init__(
+        self, model_container: ModelContainer, optimizer_container: OptimizerContainer
+    ) -> _ModelOptimizerSaver:
         self.model_container = model_container
         self.optimizer_container = optimizer_container
 

lm_engine/checkpointing/optimizer.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+from __future__ import annotations
+
 import os
 
 from torch.distributed.checkpoint.state_dict import (
@@ -15,7 +17,7 @@
 
 
 class _OptimizerSaver(Stateful):
-    def __init__(self, model_container: ModelContainer, optimizer_container: OptimizerContainer) -> None:
+    def __init__(self, model_container: ModelContainer, optimizer_container: OptimizerContainer) -> _OptimizerSaver:
         self.model_container = model_container
         self.optimizer_container = optimizer_container
 

lm_engine/containers.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+from __future__ import annotations
+
 import logging
 
 import torch.nn as nn
@@ -10,7 +12,7 @@
 
 
 class _Container:
-    def __init__(self, model_list: list[nn.Module]) -> None:
+    def __init__(self, model_list: list[nn.Module]) -> _Container:
         self.model_list = model_list
 
     def __iter__(self):
@@ -31,11 +33,11 @@ def __str__(self):
 
 
 class ModelContainer(_Container):
-    def train(self) -> "ModelContainer":
+    def train(self) -> ModelContainer:
         for model in self:
             model.train()
 
-    def eval(self) -> "ModelContainer":
+    def eval(self) -> ModelContainer:
         for model in self:
             model.eval()
 

lm_engine/data/__init__.py

@@ -89,7 +89,7 @@ def get_datasets_list(
 
 def get_finetuning_dataloader(
     args: TrainingArgs | InferenceArgs, split: DatasetSplit, mode: Mode, tokenizer: TOKENIZER_TYPE
-) -> tuple[ResumableDataLoader]:
+) -> ResumableDataLoader:
     """prepares datasets and sampler
 
     Args:
@@ -99,7 +99,7 @@ def get_finetuning_dataloader(
         tokenizer (TOKENIZER_TYPE): tokenizer
 
     Returns:
-        tuple[ResumableDataLoader]: dataloader for a blended dataset
+        ResumableDataLoader: dataloader for a blended dataset
     """
 
     assert mode == Mode.training, "blended dataset is only supported in training mode"
@@ -121,7 +121,7 @@ def get_finetuning_dataloader(
 
 def get_pretraining_dataloaders(
     args: TrainingArgs, tokenizer: TOKENIZER_TYPE, consumed_samples: int
-) -> tuple[ResumableDataLoader]:
+) -> tuple[ResumableDataLoader, list[ResumableDataLoader], list[ResumableDataLoader]]:
     if args.datasets[0].class_name == "MegatronDataset":
         dataloaders = get_megatron_gpt_dataloaders(args, tokenizer, consumed_samples=consumed_samples)
     elif args.datasets[0].class_name == "IBMDataset":
@@ -132,7 +132,7 @@ def get_pretraining_dataloaders(
 
 def _get_dispatching_dataloader(
     args: TrainingArgs | InferenceArgs, split: DatasetSplit, mode: Mode, tokenizer: TOKENIZER_TYPE
-) -> tuple[ResumableDataLoader]:
+) -> ResumableDataLoader:
     micro_batch_size = args.training_parameters.micro_batch_size
 
     num_ranks_per_node = torch.cuda.device_count()
@@ -211,7 +211,7 @@ def _get_source_broadcast_mapping() -> dict:
 
 def _get_non_dispatching_dataloader(
     args: TrainingArgs | InferenceArgs, split: DatasetSplit, mode: Mode, tokenizer: TOKENIZER_TYPE
-) -> tuple[ResumableDataLoader]:
+) -> ResumableDataLoader:
     micro_batch_size = args.training_parameters.micro_batch_size
 
     datasets_list, data_sampling_ratios = get_datasets_list(

lm_engine/data/base.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+from __future__ import annotations
+
 import torch
 
 from ..defaults import INPUT_FORMAT, OUTPUT_FORMAT
@@ -23,7 +25,7 @@ def __init__(
         output_format: str,
         max_input_tokens: int,
         max_output_tokens: int,
-    ) -> None:
+    ) -> BaseDataset:
         super().__init__()
 
         self.split = split
@@ -39,11 +41,7 @@ def __init__(
         self.do_format_output = self.output_format != OUTPUT_FORMAT
 
         # length to use for trimming (excludes eos)
-        if max_input_tokens is None:
-            self.max_input_tokens = None
-        else:
-            self.max_input_tokens = max_input_tokens
-
+        self.max_input_tokens = max_input_tokens
         self.max_output_tokens = None if max_output_tokens is None else max_output_tokens - 1
 
         self.examples = []
@@ -124,7 +122,7 @@ def __len__(self) -> int:
 class BlendedDatasets(torch.utils.data.Dataset):
     """Concatenated list of datasets for training or inference"""
 
-    def __init__(self, datasets: list[BaseDataset], split: DatasetSplit) -> None:
+    def __init__(self, datasets: list[BaseDataset], split: DatasetSplit) -> BlendedDatasets:
         super().__init__()
 
         self.split = split

lm_engine/data/dataloader.py

@@ -2,7 +2,9 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
-from typing import Callable, Iterable
+from __future__ import annotations
+
+from typing import Callable, Iterable, Iterator
 
 import torch
 import torch.distributed
@@ -37,7 +39,7 @@ def __init__(
         broadcast_world_size: int | None = None,
         static_shape_per_rank: tuple[int, int] | None = None,
         keys: list[str] = ["input_ids", "attention_mask", "labels"],
-    ) -> None:
+    ) -> DispatchingDataLoader:
         self.broadcast_world_size = broadcast_world_size
 
         self.is_source, self.source_rank, self.local_rank_in_broadcast_group, self.broadcast_group = (
@@ -67,7 +69,7 @@ def __init__(
 
         self.keys = keys
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[dict]:
         iterator = super().__iter__() if self.is_source else range(self._length)
 
         for batch in iterator:

lm_engine/data/debug.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+from __future__ import annotations
+
 from ..enums import DatasetSplit, Mode
 from ..tokenizers import TOKENIZER_TYPE
 from .base import BaseDataset
@@ -21,7 +23,7 @@ def __init__(
         output_format: str,
         max_input_tokens: int,
         max_output_tokens: int,
-    ) -> None:
+    ) -> DebugDataset:
         super().__init__(
             class_args=class_args,
             split=split,

lm_engine/data/huggingface.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+from __future__ import annotations
+
 from datasets import load_dataset
 
 from lm_engine.tokenizers import TOKENIZER_TYPE
@@ -24,7 +26,7 @@ def __init__(
         output_format: str,
         max_input_tokens: int,
         max_output_tokens: int,
-    ) -> None:
+    ) -> HuggingFaceDataset:
         super().__init__(
             class_args=class_args,
             split=split,