63 lines
1.5 KiB
Python
63 lines
1.5 KiB
Python
import argparse
|
|
|
|
def create_parser():
|
|
parser = argparse.ArgumentParser(
|
|
description="RAG pipeline CLI"
|
|
)
|
|
|
|
path_arg_parent = argparse.ArgumentParser(add_help=False)
|
|
path_arg_parent.add_argument(
|
|
"-p",
|
|
"--path",
|
|
type=str,
|
|
required=False,
|
|
help="Path to a directory containing documents to index.",
|
|
)
|
|
|
|
eval_file_arg_parent = argparse.ArgumentParser(add_help=False)
|
|
eval_file_arg_parent.add_argument(
|
|
"-f",
|
|
"--eval_file",
|
|
type=str,
|
|
required=False,
|
|
help="Path to a .json file with question/expected_answer pairs.",
|
|
)
|
|
|
|
|
|
subparsers = parser.add_subparsers(dest="commands", help="Enter a command", required=True)
|
|
|
|
subparsers.add_parser(
|
|
"run",
|
|
help="Run the full pipeline, i.e Reset, add and evaluate",
|
|
parents=[path_arg_parent, eval_file_arg_parent]
|
|
)
|
|
|
|
subparsers.add_parser(
|
|
"reset",
|
|
help="Reset the pipeline"
|
|
)
|
|
|
|
subparsers.add_parser(
|
|
"add",
|
|
help="Add the documents to the database",
|
|
parents=[path_arg_parent]
|
|
)
|
|
|
|
subparsers.add_parser(
|
|
"evaluate",
|
|
help="Evaluate the model",
|
|
parents=[eval_file_arg_parent]
|
|
)
|
|
|
|
query_parser = subparsers.add_parser(
|
|
"query",
|
|
help="Query the documents"
|
|
)
|
|
|
|
query_parser.add_argument(
|
|
"prompt",
|
|
type=str,
|
|
help="Enter the query"
|
|
)
|
|
|
|
return parser |