chore(model,prompts): add type annotations to module-level functions and __init__ methods

litgpt/model.py

@@ -930,7 +930,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 # ROPE: YaRN (Yet another RoPE extensioN) scaling function for extended context
-def yarn_get_mscale(scale=1, mscale=1):
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
     if scale <= 1:
         return 1.0
     return 0.1 * mscale * math.log(scale) + 1.0
@@ -994,7 +994,7 @@ def build_rope_cache(
 
             # Find correction range based on rotation counts
             # Inverse dimension formula to find dimension based on number of rotations
-            def find_correction_dim(num_rotations, dim, base_val, max_pos):
+            def find_correction_dim(num_rotations: float, dim: int, base_val: float, max_pos: int) -> float:
                 return (dim * math.log(max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(base_val))
 
             low_dim = find_correction_dim(beta_fast, n_elem, base, original_max_seq_len)
@@ -1072,7 +1072,7 @@ def find_correction_dim(num_rotations, dim, base_val, max_pos):
     return cos, sin
 
 
-def batched_index_select(t, dim, idx):
+def batched_index_select(t: torch.Tensor, dim: int, idx: torch.Tensor) -> torch.Tensor:
     """index_select for batched index and unbatched t"""
     if idx.dim() == 1:
         return torch.index_select(t, dim, idx)
@@ -1091,7 +1091,7 @@ def batched_index_select(t, dim, idx):
     return res
 
 
-def batched_index_copy_(t, dim, idx, val):
+def batched_index_copy_(t: torch.Tensor, dim: int, idx: torch.Tensor, val: torch.Tensor) -> torch.Tensor:
     """Index copy for batched t, idx, val"""
 
     if t.device.type == "mps":

litgpt/prompts.py

@@ -369,7 +369,7 @@ def apply(self, prompt: str, *, sys_prompt: str | None = None, **kwargs: str) ->
 
 
 class ChatML(PromptStyle):
-    def __init__(self, system_message: str | None = None):
+    def __init__(self, system_message: str | None = None) -> None:
         self.system_message = system_message
 
     def apply(self, prompt: str, *, sys_prompt: str | None = None, **kwargs: str) -> str:
@@ -380,34 +380,34 @@ def apply(self, prompt: str, *, sys_prompt: str | None = None, **kwargs: str) ->
 
 
 class Qwen2_5(ChatML):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("You are Qwen, created by Alibaba Cloud. You are a helpful assistant.")
 
 
 class Qwen2_5_Math(ChatML):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("Please reason step by step, and put your final answer within \\boxed{}.")
 
 
 class QwQ(ChatML):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             "You are a helpful and harmless assistant. You are Qwen developed by Alibaba. You should think step-by-step."
         )
 
 
 class Qwen3(ChatML):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
 
 class SmolLM2(ChatML):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("You are a helpful AI assistant named SmolLM, trained by Hugging Face")
 
 
 class Salamandra(ChatML):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             "I am Salamandra, an AI language model developed at the Barcelona Supercomputing Centre (BSC) by the Language Technologies Unit. My knowledge base was last updated on August 2023. Today Date: 2024-09-30\nSoy Salamandra, un modelo lingüístico de IA desarrollado en el Barcelona Supercomputing Centre (BSC) por la Language Technologies Unit. Mi base de conocimientos se actualizó por última vez en agosto de 2023.\nSoc Salamandra, un model de llenguatge d'IA desenvolupat al Barcelona Supercomputing Centre (BSC) per la Language Technologies Unit."
         )