CTA for leap-bundle (#3)

* CAT in readme

* CTA end of log

Author: EdoardoMosca

README.md

@@ -54,6 +54,10 @@ It uses Ray Train + Accelerate for distributed training.
 
 Unless you overwrote `output_dir`, results will be stored in `outputs/training_type/job_name/`
 
+### 3. Bundle Checkpoint for LEAP
+
+When training is done, you can bundle your output checkpoint with `leap-bundle` to use it directly within LEAP. Checkout our [Quick Start guide](https://leap.liquid.ai/docs/leap-bundle/quick-start?utm_source=github&utm_medium=link&utm_campaign=LEAP&utm_content=general).
+
 ## 🧪 Advanced Configuration
 
 ### Default Configs Location and Adding New Configs

src/leap_finetune/trainer.py

@@ -5,6 +5,8 @@
 from ray.train import RunConfig, ScalingConfig
 from ray.runtime_env import RuntimeEnv
 from ray.train.torch import TorchTrainer
+from rich.console import Console
+from rich.panel import Panel
 from torch import cuda
 
 from leap_finetune.utils.constants import RUNTIME_DIR
@@ -70,3 +72,25 @@ def ray_trainer(job_config: dict) -> None:
     )
 
     trainer.fit()
+
+    console = Console()
+    quick_start_url = (
+        "https://leap.liquid.ai/docs/leap-bundle/quick-start?utm_source=github"
+        "&utm_medium=link&utm_campaign=LEAP&utm_content=general"
+    )
+
+    cta_message = (
+        "[bold green]Training complete![/bold green]\n\n"
+        f"[bold]Checkpoint directory:[/bold] [cyan]{output_dir}[/cyan]\n\n"
+        "Bundle your output checkpoint with [bold]leap-bundle[/bold] to use it in LEAP:\n"
+        f"[dim]leap-bundle create {output_dir}[/dim]\n\n"
+        f"Quick Start: [link={quick_start_url}]{quick_start_url}[/link]"
+    )
+
+    console.print(
+        Panel.fit(
+            cta_message,
+            title="Next Step: Bundle for LEAP",
+            border_style="green",
+        )
+    )