CLI related fixes and improvements (#85)

* CLI - show image after generation w/ PIL

* CLI rename fine-tune -> fine-tuning for consistency with API

* Other CLI updates & fixes

Author: orangetin

src/together/abstract/api_requestor.py

@@ -338,7 +338,9 @@ def handle_error_response(
         try:
             assert isinstance(resp.data, dict)
             error_resp = resp.data.get("error")
-            assert isinstance(error_resp, dict)
+            assert isinstance(
+                error_resp, dict
+            ), f"Unexpected error response {error_resp}"
             error_data = TogetherErrorResponse(**(error_resp))
         except (KeyError, TypeError):
             raise error.JSONError(

src/together/cli/api/chat.py

@@ -80,11 +80,11 @@ def do_say(self, arg: str) -> None:
 
             token = chunk.choices[0].delta.content
 
-            print(token, end="", flush=True)
+            click.echo(token, nl=False)
 
             output += token
 
-        print("\n")
+        click.echo("\n")
 
         self.messages.append({"role": "assistant", "content": output})
 
@@ -239,7 +239,7 @@ def chat(
         for i, choice in enumerate(response.choices):
             if should_print_header:
                 click.echo(f"===== Completion {i} =====")
-            click.echo(json.dumps(choice.message, indent=4))
+            click.echo(choice.message.content)  # type: ignore
 
             if should_print_header:
                 click.echo("\n")

src/together/cli/api/completions.py

@@ -78,7 +78,6 @@ def completions(
                 # assertions for type checking
                 assert isinstance(stream_choice, CompletionChoicesChunk)
                 assert stream_choice.delta
-                assert stream_choice.delta.content
 
                 if should_print_header:
                     click.echo(f"\n===== Completion {stream_choice.index} =====\n")

src/together/cli/api/finetune.py

@@ -8,17 +8,17 @@
 from together.utils import finetune_price_to_dollars, parse_timestamp
 
 
-@click.group(name="fine-tune")
+@click.group(name="fine-tuning")
 @click.pass_context
-def fine_tune(ctx: click.Context) -> None:
+def fine_tuning(ctx: click.Context) -> None:
     """Fine-tunes API commands"""
     pass
 
 
-@fine_tune.command()
+@fine_tuning.command()
 @click.pass_context
 @click.option(
-    "--training_file", type=str, required=True, help="Training file ID from Files API"
+    "--training-file", type=str, required=True, help="Training file ID from Files API"
 )
 @click.option("--model", type=str, required=True, help="Base model name")
 @click.option("--n-epochs", type=int, default=1, help="Number of epochs to train for")
@@ -30,7 +30,7 @@ def fine_tune(ctx: click.Context) -> None:
 @click.option(
     "--suffix", type=str, default=None, help="Suffix for the fine-tuned model name"
 )
-@click.option("--wandb-api-key", prompt=True, hide_input=True, help="Wandb API key")
+@click.option("--wandb-api-key", type=str, default=None, help="Wandb API key")
 def create(
     ctx: click.Context,
     training_file: str,
@@ -59,10 +59,10 @@ def create(
     click.echo(json.dumps(response.model_dump(), indent=4))
 
 
-@fine_tune.command()
+@fine_tuning.command()
 @click.pass_context
 def list(ctx: click.Context) -> None:
-    """List fine-tuning tasks"""
+    """List fine-tuning jobs"""
     client: Together = ctx.obj
 
     response = client.fine_tuning.list()
@@ -89,38 +89,34 @@ def list(ctx: click.Context) -> None:
     click.echo(table)
 
 
-@fine_tune.command()
+@fine_tuning.command()
 @click.pass_context
 @click.argument("fine_tune_id", type=str, required=True)
 def retrieve(ctx: click.Context, fine_tune_id: str) -> None:
-    """Retrieve fine-tuning task"""
+    """Retrieve fine-tuning job details"""
     client: Together = ctx.obj
 
     response = client.fine_tuning.retrieve(fine_tune_id)
 
-    table_data = [
-        {"Key": key, "Value": value}
-        for key, value in response.model_dump().items()
-        if key not in ["events"]
-    ]
-    table = tabulate(table_data, tablefmt="grid")
+    # remove events from response for cleaner output
+    response.events = None
 
-    click.echo(table)
+    click.echo(json.dumps(response.model_dump(), indent=4))
 
 
-@fine_tune.command()
+@fine_tuning.command()
 @click.pass_context
 @click.argument("fine_tune_id", type=str, required=True)
 def cancel(ctx: click.Context, fine_tune_id: str) -> None:
-    """Cancel fine-tuning task"""
+    """Cancel fine-tuning job"""
     client: Together = ctx.obj
 
     response = client.fine_tuning.cancel(fine_tune_id)
 
     click.echo(json.dumps(response.model_dump(), indent=4))
 
 
-@fine_tune.command()
+@fine_tuning.command()
 @click.pass_context
 @click.argument("fine_tune_id", type=str, required=True)
 def list_events(ctx: click.Context, fine_tune_id: str) -> None:
@@ -135,7 +131,7 @@ def list_events(ctx: click.Context, fine_tune_id: str) -> None:
     for i in response.data:
         display_list.append(
             {
-                "Message": i.message,
+                "Message": "\n".join(wrap(i.message or "", width=50)),
                 "Type": i.type,
                 "Created At": parse_timestamp(i.created_at or ""),
                 "Hash": i.hash,
@@ -146,14 +142,14 @@ def list_events(ctx: click.Context, fine_tune_id: str) -> None:
     click.echo(table)
 
 
-@fine_tune.command()
+@fine_tuning.command()
 @click.pass_context
 @click.argument("fine_tune_id", type=str, required=True)
 @click.option(
     "--output_dir",
     type=click.Path(exists=True, file_okay=False, resolve_path=True),
     required=False,
-    default=".",
+    default=None,
     help="Output directory",
 )
 @click.option(

src/together/cli/api/images.py

@@ -2,6 +2,7 @@
 import pathlib
 
 import click
+from PIL import Image
 
 from together import Together
 from together.types import ImageResponse
@@ -33,6 +34,7 @@ def images(ctx: click.Context) -> None:
     help="Output directory",
 )
 @click.option("--prefix", type=str, required=False, default="image-")
+@click.option("--no-show", is_flag=True, help="Do not open images in viewer")
 def generate(
     ctx: click.Context,
     prompt: str,
@@ -45,6 +47,7 @@ def generate(
     negative_prompt: str,
     output: pathlib.Path,
     prefix: str,
+    no_show: bool,
 ) -> None:
     """Generate image"""
 
@@ -63,8 +66,17 @@ def generate(
 
     assert isinstance(response, ImageResponse)
     assert isinstance(response.data, list)
-    for choice in response.data:
+
+    for i, choice in enumerate(response.data):
         assert isinstance(choice, ImageChoicesData)
 
         with open(f"{output}/{prefix}{choice.index}.png", "wb") as f:
             f.write(base64.b64decode(choice.b64_json))
+
+        click.echo(
+            f"Image [{i + 1}/{len(response.data)}] saved to {output}/{prefix}{choice.index}.png"
+        )
+
+        if not no_show:
+            image = Image.open(f"{output}/{prefix}{choice.index}.png")
+            image.show()

src/together/cli/cli.py

@@ -9,7 +9,7 @@
 from together.cli.api.chat import chat, interactive
 from together.cli.api.completions import completions
 from together.cli.api.files import files
-from together.cli.api.finetune import fine_tune
+from together.cli.api.finetune import fine_tuning
 from together.cli.api.images import images
 from together.cli.api.models import models
 from together.constants import MAX_RETRIES, TIMEOUT_SECS
@@ -70,7 +70,7 @@ def main(
 main.add_command(completions)
 main.add_command(images)
 main.add_command(files)
-main.add_command(fine_tune)
+main.add_command(fine_tuning)
 main.add_command(models)
 
 if __name__ == "__main__":

src/together/filemanager.py

@@ -36,19 +36,19 @@ def chmod_and_replace(src: Path, dst: Path) -> None:
     """
 
     # Get umask by creating a temporary file in the cache folder.
-    tmp_file = dst.parent.parent / f"tmp_{uuid.uuid4()}"
+    tmp_file = dst.parent / f"tmp_{uuid.uuid4()}"
 
     try:
         tmp_file.touch()
 
         cache_dir_mode = Path(tmp_file).stat().st_mode
 
-        os.chmod(src, stat.S_IMODE(cache_dir_mode))
+        os.chmod(src.as_posix(), stat.S_IMODE(cache_dir_mode))
 
     finally:
         tmp_file.unlink()
 
-    shutil.move(src, dst)
+    shutil.move(src.as_posix(), dst.as_posix())
 
 
 def _get_file_size(
@@ -237,7 +237,7 @@ def download(
 
         os.remove(lock_path)
 
-        return file_path.as_posix(), file_size
+        return str(file_path.resolve()), file_size
 
 
 class UploadManager:

src/together/resources/finetune.py

@@ -219,7 +219,7 @@ def download(
             output = Path(output)
 
         downloaded_filename, file_size = download_manager.download(
-            url, output, normalize_key(remote_name or ""), fetch_metadata=True
+            url, output, normalize_key(remote_name or id), fetch_metadata=True
         )
 
         return FinetuneDownloadResult(

src/together/types/finetune.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from enum import Enum
-from pathlib import Path
 from typing import List, Literal
 
 from pydantic import Field
@@ -202,6 +201,6 @@ class FinetuneDownloadResult(BaseModel):
     # checkpoint step number
     checkpoint_step: int | None = None
     # local path filename
-    filename: Path | None = None
+    filename: str | None = None
     # size in bytes
     size: int | None = None