Added a "--raw" option for model template printing

This commit is contained in:
Ashwin Bharambe 2024-09-03 22:04:43 -07:00
parent 85d56ed3f2
commit 86059af5af

View file

@ -59,10 +59,15 @@ class ModelTemplate(Subcommand):
self.parser.add_argument( self.parser.add_argument(
"--format", "--format",
type=str, 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, required=False,
default="json", 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: def _run_model_template_cmd(self, args: argparse.Namespace) -> None:
from llama_models.llama3.api.interface import ( from llama_models.llama3.api.interface import (
@ -82,15 +87,23 @@ class ModelTemplate(Subcommand):
else: else:
rendered += tok rendered += tok
rendered = rendered.replace("\n", "\n") if not args.raw:
print_table( rendered = rendered.replace("\n", "\n")
[ print_table(
("Name", colored(template.template_name, "white", attrs=["bold"])), [
("Template", rendered), (
("Notes", template.notes), "Name",
], colored(template.template_name, "white", attrs=["bold"]),
separate_rows=True, ),
) ("Template", rendered),
("Notes", template.notes),
],
separate_rows=True,
)
else:
print("Template: ", template.template_name)
print("=" * 40)
print(rendered)
else: else:
templates = list_jinja_templates() templates = list_jinja_templates()
headers = ["Role", "Template Name"] headers = ["Role", "Template Name"]