Remove hardcoded combined.jsonl with a flag

mosaicml · Jan 12, 2024 · 876fab5 · 876fab5
1 parent fa8f3d9
commit 876fab5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/scripts/data_prep/convert_delta_to_json.py b/scripts/data_prep/convert_delta_to_json.py
@@ -460,7 +460,7 @@ def fetch_DT(args: Namespace) -> None:
     # combine downloaded jsonl into one big jsonl for IFT
     iterative_combine_jsons(
         args.json_output_path,
-        os.path.join(args.json_output_path, 'combined.jsonl'))
+        os.path.join(args.json_output_path, args.json_output_filename))
 
 
 if __name__ == '__main__':
@@ -505,6 +505,12 @@ def fetch_DT(args: Namespace) -> None:
         help=
         'Use serverless or not. Make sure the workspace is entitled with serverless'
     )
+    parser.add_argument(
+        '--json_output_filename',
+        required=False,
+        type=str,
+        default='train-00000-of-00001.jsonl',
+        help='The combined final jsonl that combines all partitioned jsonl')
     args = parser.parse_args()
 
     from databricks.sdk import WorkspaceClient

diff --git a/tests/a_scripts/data_prep/test_convert_delta_to_json.py b/tests/a_scripts/data_prep/test_convert_delta_to_json.py
@@ -36,6 +36,7 @@ def test_stream_delta_to_json(self, mock_workspace_client: Any,
         args.cluster_id = '1234'
         args.debug = False
         args.use_serverless = False
+        args.json_output_filename = 'combined.jsonl'
 
         mock_cluster_get = MagicMock()
         mock_cluster_get.return_value = MagicMock(