Merge pull request #27 from rajveer43/patch-1

add docstrings

Author: ctlllll

medusa/model/kv_cache.py

@@ -4,6 +4,7 @@
 class KVCache:
     """
     A key-value cache for the model.
+
     This class provides a mechanism to maintain a growing cache of keys and values,
     particularly useful for models that benefit from caching previous states,
     like transformers during autoregressive decoding.
@@ -15,6 +16,8 @@ class KVCache:
 
     def __init__(self, data, current_length):
         """
+        Initialize the KVCache.
+
         Args:
             data (torch.Tensor): Initial tensor to store the keys and values.
             current_length (int): Initial length of the data.

medusa/model/medusa_model.py

@@ -10,6 +10,16 @@
 
 
 class MedusaConfig(PretrainedConfig):
+    """
+    Configuration class for Medusa model.
+
+    Args:
+        medusa_num_heads (int, optional): Number of heads for the Medusa layer. Default is 2.
+        medusa_num_layers (int, optional): Number of Medusa layers. Default is 1.
+        base_model_name_or_path (str, optional): The name or path of the base model. Default is "lmsys/vicuna-7b-v1.3".
+        **kwargs: Additional keyword arguments to be passed to the parent class constructor.
+    """
+
     def __init__(
         self,
         medusa_num_heads=2,
@@ -24,10 +34,14 @@ def __init__(
 
 
 class ResBlock(nn.Module):
-    """A Residual Block module.
+    """
+    A Residual Block module.
 
     This module performs a linear transformation followed by a SiLU activation,
     and then adds the result to the original input, creating a residual connection.
+
+    Args:
+        hidden_size (int): The size of the hidden layers in the block.
     """
 
     def __init__(self, hidden_size):
@@ -39,7 +53,8 @@ def __init__(self, hidden_size):
         self.act = nn.SiLU()
 
     def forward(self, x):
-        """Forward pass of the ResBlock.
+        """
+        Forward pass of the ResBlock.
 
         Args:
             x (torch.Tensor): Input tensor.