Fix pyright warnings/errors

This patch fixes a pyright string concatenation warning and also adds typing information where necessary.
mosaicml · Sep 20, 2023 · 1831916 · 1831916
1 parent 83d36a5
commit 1831916
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/llmfoundry/data/finetuning/tasks.py b/llmfoundry/data/finetuning/tasks.py
@@ -206,7 +206,7 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
             return _read_binary_tokenized_sample(sample)
         else:
             raise RuntimeError(
-                'FineTurningDataset needs samples to have prompt/response columns '
+                'FineTurningDataset needs samples to have prompt/response columns ' +\
                 'or tokens/labels columns')
 
 

diff --git a/tests/test_dataloader.py b/tests/test_dataloader.py
@@ -63,7 +63,7 @@ def build_mock_ft_streaming_dataset(data_path: str, split: str):
             output_writer.write(sample)
 
 
-def build_mock_tokenized_ft_streaming_dataset(data_path, split):
+def build_mock_tokenized_ft_streaming_dataset(data_path: str, split: str):
     columns = {'tokens': 'bytes', 'labels': 'bytes'}
 
     dataset = [{