shorten docstring

src/delphi/eval/token_map.py

@@ -9,17 +9,8 @@ def token_map(
     tokenized_dataset: Dataset,
     output_path: str | None = None,
 ) -> dict[int, list[tuple[int, int]]]:
-    """Return a mapping of tokens to their (prompt_idx, token_idx) locations in the tokenized_dataset.
+    """Return a mapping of tokens to their (prompt_idx, token_idx) locations in the tokenized_dataset."""
 
-    Args:
-        tokenized_dataset (Dataset): A tokenized dataset.
-        save_output (bool, optional): Whether to save the output to a file. Defaults to True.
-        output_path (str, optional): The output file path. Defaults to "/data/token_map.pkl".
-
-    Returns:
-        dict[int, list[tuple[int, int]]]: A mapping of tokens to their (prompt_idx, token_idx)
-            locations in the tokenized_dataset.
-    """
     mapping = {}
     for prompt_idx, prompt in enumerate(tokenized_dataset):
         prompt = cast(dict, prompt)