diff --git a/src/together/lib/cli/__init__.py b/src/together/lib/cli/__init__.py index 0561ac7d..a64acd8b 100644 --- a/src/together/lib/cli/__init__.py +++ b/src/together/lib/cli/__init__.py @@ -71,6 +71,7 @@ BETA_CLUSTERS_UPDATE_HELP_EXAMPLES, BETA_MODELS_DOWNLOAD_HELP_EXAMPLES, FINE_TUNING_DOWNLOAD_HELP_EXAMPLES, + FINE_TUNING_PREVIEW_HELP_EXAMPLES, BETA_CLUSTERS_STORAGE_HELP_EXAMPLES, BETA_ENDPOINTS_DEPLOY_HELP_EXAMPLES, BETA_ENDPOINTS_SHADOW_HELP_EXAMPLES, @@ -454,6 +455,11 @@ async def run_command() -> None: help="Retrieve training metrics for a fine-tuning job", help_epilogue=FINE_TUNING_LIST_METRICS_HELP_EXAMPLES, ) +fine_tuning_app.command( + (f"{_CLI}.fine_tuning.preview:preview"), + help="Preview how a training file will be tokenized", + help_epilogue=FINE_TUNING_PREVIEW_HELP_EXAMPLES, +) ## Models API commands models_app = app.command(App(name="models", help="List and upload models", help_epilogue=MODELS_HELP_EXAMPLES)) diff --git a/src/together/lib/cli/api/fine_tuning/preview.py b/src/together/lib/cli/api/fine_tuning/preview.py new file mode 100644 index 00000000..671165bc --- /dev/null +++ b/src/together/lib/cli/api/fine_tuning/preview.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +from typing import Optional, Annotated +from typing_extensions import Literal + +from cyclopts import Parameter +from rich.markup import escape as escape_rich_markup + +from together import omit +from together._utils._json import openapi_dumps +from together.lib.cli.utils.config import CLIConfigParameter +from together.lib.cli.utils._console import console +from together.lib.cli.components.list import ListTable +from together.lib.cli.components.loader import show_loading_status + +_TOKEN_PREVIEW_LIMIT = 24 + + +def _format_spans(spans: list[list[int]]) -> str: + if len(spans) == 0: + return "-" + + formatted = [] + for span in spans: + if len(span) == 2: + formatted.append(f"{span[0]}:{span[1]}") + return ", ".join(formatted) if formatted else "-" + + +def _format_tokens(tokens: list[str], labels: list[int]) -> str: + if len(tokens) == 0: + return "-" + + rendered_tokens: list[str] = [] + for index, token in enumerate(tokens[:_TOKEN_PREVIEW_LIMIT]): + text = escape_rich_markup(repr(token)) + if index < len(labels) and labels[index] != -100: + rendered_tokens.append(f"[primary]{text}[/primary]") + else: + rendered_tokens.append(f"[dim]{text}[/dim]") + + remaining = len(tokens) - _TOKEN_PREVIEW_LIMIT + if remaining > 0: + rendered_tokens.append(f"[dim]... +{remaining} tokens[/dim]") + + return " ".join(rendered_tokens) + + +async def preview( + training_file: Annotated[ + str, + Parameter( + alias="-t", + help="Training file ID from the Files API to sample for preview", + ), + ], + model: Annotated[ + str, + Parameter( + alias="-M", + help="Name of the base model whose tokenizer and chat template will be used", + ), + ], + *, + config: CLIConfigParameter, + top_k: Annotated[ + Optional[int], + Parameter(help="Maximum number of rows from the start of the training file to tokenize"), + ] = None, + train_on_inputs: Annotated[ + Optional[bool], + Parameter( + negative=(), + help="Whether prompt or user-message tokens should contribute to training loss", + ), + ] = None, + training_method: Annotated[ + Literal["sft"], + Parameter( + alias="-m", + help="Fine-tuning method to preview. Only supervised fine-tuning is currently supported.", + ), + ] = "sft", +) -> None: + """Preview how a fine-tuning training file will be tokenized.""" + response = await show_loading_status( + "Generating fine-tuning preview...", + config.client.fine_tuning.preview( + model=model, + training_file=training_file, + top_k=top_k if top_k is not None else omit, + train_on_inputs=train_on_inputs if train_on_inputs is not None else omit, + training_method=training_method, + ), + ) + + if config.json: + console.print_json(openapi_dumps(response).decode("utf-8")) + return + + console.print(f"[muted]Model:[/muted] {response.model}") + console.print(f"[muted]Dataset format:[/muted] {response.dataset_format}") + console.print(f"[muted]Max sequence length:[/muted] {response.max_seq_length}") + console.print(f"[muted]Train on inputs:[/muted] {response.train_on_inputs}") + + table = ListTable(title="Preview rows", empty_message="No preview rows returned") + table.add_primary_column("Row") + table.add_column("Tokens", ratio=4) + table.add_column("Trained Spans") + table.add_column("Trained / Total", justify="right") + table.add_column("Truncated", justify="center") + + for index, row in enumerate(response.rows, start=1): + table.add_row( + str(index), + _format_tokens(row.tokens, row.labels), + _format_spans(row.trained_spans), + f"{row.num_trained_tokens} / {row.num_tokens}", + "yes" if row.truncated else "no", + ) + + console.print(table) diff --git a/src/together/lib/cli/utils/_help_examples.py b/src/together/lib/cli/utils/_help_examples.py index ec5e7bc0..780882dc 100644 --- a/src/together/lib/cli/utils/_help_examples.py +++ b/src/together/lib/cli/utils/_help_examples.py @@ -91,6 +91,9 @@ [dim]-[/dim] Cancel a fine-tuning job: [primary]tg ft cancel [/primary] + +[dim]-[/dim] Preview how a training file will be tokenized: + [primary]tg ft preview --model Qwen/Qwen2-1.5B --training-file [/primary] """ FINE_TUNING_CREATE_HELP_EXAMPLES = """[dim]Examples:[/dim] @@ -127,6 +130,17 @@ [primary]tg ft list-metrics > plots.txt[/primary] """ +FINE_TUNING_PREVIEW_HELP_EXAMPLES = """[dim]Examples:[/dim] +[dim]-[/dim] Preview tokenization for the first rows in a training file: + [primary]tg ft preview --model Qwen/Qwen2-1.5B --training-file [/primary] + +[dim]-[/dim] Preview more rows and include prompt tokens in training loss: + [primary]tg ft preview -M Qwen/Qwen2-1.5B -t --top-k 10 --train-on-inputs[/primary] + +[dim]-[/dim] Save the raw preview response: + [primary]tg ft preview -M Qwen/Qwen2-1.5B -t --json > preview.json[/primary] +""" + FINE_TUNING_DOWNLOAD_HELP_EXAMPLES = """[dim]Examples:[/dim] [dim]-[/dim] Download a fine-tuned model's weights: [primary]tg ft download --output-dir ./my-model[/primary] diff --git a/tests/cli/test_fine_tuning.py b/tests/cli/test_fine_tuning.py index 5a17a6fc..d96a161c 100644 --- a/tests/cli/test_fine_tuning.py +++ b/tests/cli/test_fine_tuning.py @@ -96,6 +96,24 @@ "status": "pending", } +_FT_PREVIEW_BODY = { + "dataset_format": "instruction", + "max_seq_length": 4096, + "model": "meta-llama/Llama-3-8b", + "train_on_inputs": True, + "rows": [ + { + "input_ids": [101, 102, 103], + "labels": [-100, 102, 103], + "num_tokens": 3, + "num_trained_tokens": 2, + "tokens": ["", "hello", "world"], + "trained_spans": [[1, 3]], + "truncated": False, + } + ], +} + class TestFineTuningCreate: @pytest.mark.respx(base_url=base_url) @@ -269,6 +287,61 @@ def test_ft_alias_list(self, respx_mock: MockRouter, cli_runner: CliRunner) -> N assert "ft-newer" in result.output +class TestFineTuningPreview: + @pytest.mark.respx(base_url=base_url) + def test_preview_json_sends_params(self, respx_mock: MockRouter, cli_runner: CliRunner) -> None: + route = respx_mock.post("/fine-tunes/preview").mock(return_value=httpx.Response(200, json=_FT_PREVIEW_BODY)) + + result = cli_runner.invoke( + [ + "fine-tuning", + "preview", + "--training-file", + "file-train", + "--model", + "meta-llama/Llama-3-8b", + "--top-k", + "2", + "--train-on-inputs", + "--json", + ] + ) + + assert result.exit_code == 0 + body = json.loads(cast(Call, route.calls[0]).request.content) + assert body == { + "model": "meta-llama/Llama-3-8b", + "training_file": "file-train", + "top_k": 2, + "train_on_inputs": True, + "training_method": "sft", + } + parsed = json.loads(result.output) + assert parsed["dataset_format"] == "instruction" + assert parsed["rows"][0]["tokens"] == ["", "hello", "world"] + + @pytest.mark.respx(base_url=base_url) + def test_preview_table(self, respx_mock: MockRouter, cli_runner: CliRunner) -> None: + respx_mock.post("/fine-tunes/preview").mock(return_value=httpx.Response(200, json=_FT_PREVIEW_BODY)) + + result = cli_runner.invoke( + [ + "ft", + "preview", + "-t", + "file-train", + "-M", + "meta-llama/Llama-3-8b", + ] + ) + + assert result.exit_code == 0 + assert "Dataset format" in result.output + assert "instruction" in result.output + assert "hello" in result.output + assert "2 / 3" in result.output + + class TestFineTuningRetrieve: @pytest.mark.respx(base_url=base_url) def test_retrieve_json(self, respx_mock: MockRouter, cli_runner: CliRunner) -> None: