Set default output name for model weights and file downloads

Author: orangetin

src/together/commands/files.py

@@ -68,10 +68,13 @@ def _add_retrieve_content(
         type=str,
     )
     retrieve_file_content_parser.add_argument(
-        "output",
+        "--output",
+        "-o",
+        default=None,
         metavar="OUT_FILENAME",
         help="Output filename",
         type=str,
+        required=False,
     )
 
     retrieve_file_content_parser.set_defaults(func=_run_retrieve_content)

src/together/commands/finetune.py

@@ -175,10 +175,12 @@ def _add_download(
         type=str,
     )
     download_parser.add_argument(
-        "output",
-        metavar="OUT_FILENAME",
+        "--output",
+        "-o",
+        default=None,
         help="Output filename",
         type=str,
+        required=False,
     )
     download_parser.add_argument(
         "--checkpoint-num",

src/together/files.py

@@ -129,7 +129,12 @@ def retrieve_file(self, file_id: str) -> Dict[str, Union[str, int]]:
 
         return response_json
 
-    def retrieve_file_content(self, file_id: str, output: str) -> str:
+    def retrieve_file_content(
+        self, file_id: str, output: Union[str, None] = None
+    ) -> str:
+        if output is None:
+            output = file_id + ".jsonl"
+
         relative_path = posixpath.join(file_id, "content")
         retrieve_url = urllib.parse.urljoin(self.endpoint_url, relative_path)
 

src/together/finetune.py

@@ -1,7 +1,7 @@
 import os
 import posixpath
 import urllib.parse
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import requests
 from tqdm import tqdm
@@ -206,7 +206,19 @@ def is_final_model_available(self, fine_tune_id: str) -> bool:
                     return True
         return False
 
-    def download(self, fine_tune_id: str, output: str, checkpoint_num: int = -1) -> str:
+    def download(
+        self,
+        fine_tune_id: str,
+        output: Union[str, None] = None,
+        checkpoint_num: int = -1,
+    ) -> str:
+        # default to model_output_path name
+        if output is None:
+            output = (
+                self.retrieve_finetune(fine_tune_id)["model_output_path"].split("/")[-1]
+                + ".tar.gz"
+            )
+
         model_file_path = urllib.parse.urljoin(
             self.endpoint_url,
             f"/api/finetune/downloadfinetunefile?ft_id={fine_tune_id}",