add help and descriptions to cli

This commit is contained in:
Bruce MacDonald 2023-06-30 12:14:59 -04:00
parent 36168300d9
commit ab49a18a33

View file

@ -1,32 +1,76 @@
import os import os
import sys import sys
from argparse import ArgumentParser from argparse import ArgumentParser, HelpFormatter, PARSER
from yaspin import yaspin from yaspin import yaspin
from ollama import model, engine from ollama import model, engine
from ollama.cmd import server from ollama.cmd import server
class CustomHelpFormatter(HelpFormatter):
"""
This class is used to customize the way the argparse help text is displayed.
We specifically override the _format_action method to exclude the line that
shows all the subparser command options in the help text. This line is typically
in the form "{serve,models,pull,run}".
"""
def _format_action(self, action):
# get the original help text
parts = super()._format_action(action)
if action.nargs == PARSER:
# remove the unwanted first line
parts = "\n".join(parts.split("\n")[1:])
return parts
def main(): def main():
parser = ArgumentParser() parser = ArgumentParser(
description='Ollama: Run any large language model on any machine.',
formatter_class=CustomHelpFormatter,
)
# create models home if it doesn't exist # create models home if it doesn't exist
os.makedirs(model.models_home, exist_ok=True) os.makedirs(model.models_home, exist_ok=True)
subparsers = parser.add_subparsers() subparsers = parser.add_subparsers(
title='commands',
)
server.set_parser(subparsers.add_parser("serve")) server.set_parser(
subparsers.add_parser(
"serve",
description="Start a persistent server to interact with models via the API.",
help="Start a persistent server to interact with models via the API.",
)
)
list_parser = subparsers.add_parser("models") list_parser = subparsers.add_parser(
"models",
description="List all available models stored locally.",
help="List all available models stored locally.",
)
list_parser.set_defaults(fn=list_models) list_parser.set_defaults(fn=list_models)
pull_parser = subparsers.add_parser("pull") pull_parser = subparsers.add_parser(
pull_parser.add_argument("model") "pull",
description="Download a specified model from a remote source.",
help="Download a specified model from a remote source. Usage: pull [model]",
)
pull_parser.add_argument("model", help="Name of the model to download.")
pull_parser.set_defaults(fn=pull) pull_parser.set_defaults(fn=pull)
run_parser = subparsers.add_parser("run") run_parser = subparsers.add_parser(
run_parser.add_argument("model") "run",
run_parser.add_argument("prompt", nargs="?") description="Run a model and submit prompts.",
help="Run a model and submit prompts. Usage: run [model] [prompt]",
)
run_parser.add_argument("model", help="Name of the model to run.")
run_parser.add_argument(
"prompt",
nargs="?",
help="Optional prompt for the model, interactive mode enabled when not specified.",
)
run_parser.set_defaults(fn=run) run_parser.set_defaults(fn=run)
args = parser.parse_args() args = parser.parse_args()