From 86059af5af4483fe3c43a42621e980556cae6664 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Tue, 3 Sep 2024 22:04:43 -0700
Subject: [PATCH] Added a "--raw" option for model template printing

---
 llama_toolchain/cli/model/template.py | 33 +++++++++++++++++++--------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/llama_toolchain/cli/model/template.py b/llama_toolchain/cli/model/template.py
index ca898e7a6..2776d9703 100644
--- a/llama_toolchain/cli/model/template.py
+++ b/llama_toolchain/cli/model/template.py
@@ -59,10 +59,15 @@ class ModelTemplate(Subcommand):
         self.parser.add_argument(
             "--format",
             type=str,
-            help="ToolPromptFormat ( json or functino_tag). This flag is used to print the template in a specific formats.",
+            help="ToolPromptFormat (json or function_tag). This flag is used to print the template in a specific formats.",
             required=False,
             default="json",
         )
+        self.parser.add_argument(
+            "--raw",
+            action="store_true",
+            help="If set to true, don't pretty-print into a table. Useful to copy-paste.",
+        )
 
     def _run_model_template_cmd(self, args: argparse.Namespace) -> None:
         from llama_models.llama3.api.interface import (
@@ -82,15 +87,23 @@ class ModelTemplate(Subcommand):
                 else:
                     rendered += tok
 
-            rendered = rendered.replace("\n", "↵\n")
-            print_table(
-                [
-                    ("Name", colored(template.template_name, "white", attrs=["bold"])),
-                    ("Template", rendered),
-                    ("Notes", template.notes),
-                ],
-                separate_rows=True,
-            )
+            if not args.raw:
+                rendered = rendered.replace("\n", "↵\n")
+                print_table(
+                    [
+                        (
+                            "Name",
+                            colored(template.template_name, "white", attrs=["bold"]),
+                        ),
+                        ("Template", rendered),
+                        ("Notes", template.notes),
+                    ],
+                    separate_rows=True,
+                )
+            else:
+                print("Template: ", template.template_name)
+                print("=" * 40)
+                print(rendered)
         else:
             templates = list_jinja_templates()
             headers = ["Role", "Template Name"]