Add doc strings

utils.py

@@ -79,7 +79,50 @@ def get_gpu(gpu: DEVICE_TYPE) -> ENCODER_DEVICE_TYPE:
 
 
 def slice_embeddings(embeddings: NDArray, num_sentences: NumSentencesType) -> EmbeddingSlicesType:
+    """
+        Slice embeddings into segments based on the provided number of sentences per segment.
+
+        Args:
+        - embeddings (np.ndarray): The array of embeddings to be sliced.
+        - num_sentences (Union[List[int], List[List[int]]]):
+            - If a list of integers: Specifies the number of embeddings to take in each slice.
+            - If a list of lists of integers: Specifies multiple nested levels of slicing.
+
+        Returns:
+        - List[np.ndarray]: A list of numpy arrays where each array represents a slice of embeddings.
+
+        Raises:
+        - TypeError: If `num_sentences` is not of type List[int] or List[List[int]].
+
+        Example Usage:
+
+        ```python
+        embeddings = np.random.rand(10, 5)
+        num_sentences = [3, 2, 5]
+        result = slice_embeddings(embeddings, num_sentences)
+        # `result` will be a list of numpy arrays:
+        # [embeddings[:3], embeddings[3:5], embeddings[5:]]
+
+        num_sentences_nested = [[2, 1], [3, 4]]
+        result_nested = slice_embeddings(embeddings, num_sentences_nested)
+        # `result_nested` will be a nested list of numpy arrays:
+        # [[embeddings[:2], embeddings[2:3]], [embeddings[3:6], embeddings[6:]]]
+
+        slice_embeddings(embeddings, "invalid")  # Raises a TypeError
+        ```
+    """
+
     def _slice_embeddings(s_idx: int, n_sentences: List[int]):
+        """
+            Helper function to slice embeddings starting from index `s_idx`.
+
+            Args:
+            - s_idx (int): Starting index for slicing.
+            - n_sentences (List[int]): List specifying number of sentences in each slice.
+
+            Returns:
+            - Tuple[List[np.ndarray], int]: A tuple containing a list of sliced embeddings and the next starting index.
+        """
         _result = []
         for count in n_sentences:
             _result.append(embeddings[s_idx:s_idx + count])
@@ -107,6 +150,37 @@ def slice_embeddings(embeddings: NDArray, num_sentences: NumSentencesType) -> Em
 
 
 def is_nested_list_of_type(lst_obj, element_type, depth: int) -> bool:
+    """
+        Check if the given object is a nested list of a specific type up to a specified depth.
+
+        Args:
+        - lst_obj: The object to check, expected to be a list or a single element.
+        - element_type: The type that each element in the nested list should match.
+        - depth (int): The depth of nesting to check. Must be non-negative.
+
+        Returns:
+        - bool: True if lst_obj is a nested list of the specified type up to the given depth, False otherwise.
+
+        Raises:
+        - ValueError: If depth is negative.
+
+        Example:
+        ```python
+        # Test cases
+        is_nested_list_of_type("test", str, 0)  # Returns True
+        is_nested_list_of_type([1, 2, 3], str, 0)  # Returns False
+        is_nested_list_of_type(["apple", "banana"], str, 1)  # Returns True
+        is_nested_list_of_type([[1, 2], [3, 4]], int, 2)  # Returns True
+        is_nested_list_of_type([[1, 2], ["a", "b"]], int, 2)  # Returns False
+        is_nested_list_of_type([[[1], [2]], [[3], [4]]], int, 3)  # Returns True
+        ```
+
+        Explanation:
+        - The function checks if `lst_obj` is a nested list of elements of type `element_type` up to `depth` levels deep.
+        - If `depth` is 0, it checks if `lst_obj` itself is of type `element_type`.
+        - If `depth` is greater than 0, it recursively checks each level of nesting to ensure all elements match `element_type`.
+        - Raises a `ValueError` if `depth` is negative, as depth must be a non-negative integer.
+    """
     if depth == 0:
         return isinstance(lst_obj, element_type)
     elif depth > 0: