def parse_args()

in assets/aml-benchmark/components/src/aml_benchmark/prompt_crafter/main.py [0:0]

• 86 lines of code
• 1 McCabe index (conditional complexity)

def parse_args() -> ArgumentParser:
    """Parse command line arguments."""
    parser = ArgumentParser()
    parser.add_argument(
        "--test_data",
        type=str,
        required=True,
        help="Path to load the test data to generate prompts."
        )
    parser.add_argument(
        "--few_shot_data",
        type=str,
        required=False,
        help="The uri file(jsonl) to generate n-shot prompts.")
    parser.add_argument(
        "--prompt_type",
        type=str,
        required=True,
        choices=[PromptType.chat.name, PromptType.completions.name],
        help="The type of prompt to be generated.")
    parser.add_argument(
        "--n_shots",
        type=int,
        required=True,
        help="The number of shots to use for n-shot prompts.")
    parser.add_argument(
        "--random_seed",
        type=int,
        default=0,
        required=False,
        help="The random seed to used for generating prompts.")
    parser.add_argument(
        "--output_file",
        type=str,
        required=True,
        help="Path to output jsonl which has the generated prompts.")
    parser.add_argument(
        "--prompt_pattern",
        type=str,
        required=True,
        help="The jinja pattern to used for generating prompts.")
    parser.add_argument(
        "--few_shot_separator",
        type=str,
        default="",
        required=False,
        help="The separator to be used for n-shot prompts.")
    parser.add_argument(
        "--prefix",
        type=str,
        required=False,
        help="The prefix to be used for prompts.")
    parser.add_argument(
        "--label_map",
        type=str,
        required=False,
        help="The label map to be used for prompts.")
    parser.add_argument(
        "--ground_truth_column_name",
        type=str,
        required=False,
        help="The ground truth key in the input data.")
    parser.add_argument(
        "--additional_columns",
        type=str,
        required=False,
        help="The additional columns in the input data, seperated by comma.")
    parser.add_argument(
        "--system_message",
        type=str,
        required=False,
        help="The system message to be used for chat prompts.")
    parser.add_argument(
        "--output_pattern",
        type=str,
        required=True,
        help="The jinja template representing the expected output to \
         be used for few shot prompts when n_shot > 0")
    parser.add_argument(
        "--few_shot_pattern",
        type=str,
        required=False,
        default=None,
        help="The jinja template used to generate few shot prompts. \
            Few shot prompts use the pattern in the `prompt_pattern` input \
            if `few_shot_pattern` is not provided.")
    return parser.parse_args()