allow for custom prompts

Author: gabriben

src/wraval/actions/action_llm_judge.py

@@ -7,6 +7,7 @@
 from dynaconf import Dynaconf
 from .data_utils import write_dataset_local, write_dataset_to_s3, load_latest_dataset
 from .prompts_judge import generate_input_prompt, generate_system_prompt, get_rubric, rewrite_prompt
+
 from .completion import batch_get_bedrock_completions
 import re
 import boto3
@@ -58,6 +59,10 @@ def process_tone_data(
     Returns:
         Processed DataFrame with scores
     """
+
+    if settings.custom_prompts == True:
+        from wraval.custom_prompts.prompts_judge import generate_input_prompt, generate_system_prompt
+
     dmt = d[d.tone == tone].copy()
     rubrics = list(tone_rubrics.keys())
 
@@ -111,6 +116,10 @@ def judge(
         data_dir: Directory containing input data
         endpoint_type: Type of endpoint to use
     """
+
+    if settings.custom_prompts == True:
+        from wraval.custom_prompts.prompts_judge import get_rubric
+
     try:
         d = load_latest_dataset(settings.data_dir)
         print(f"Loaded dataset with {len(d)} rows")

src/wraval/main.py

@@ -32,6 +32,11 @@ def get_settings(args):
     ## add the AWS account you are logged into, if necessary.
     settings.model = settings.model.format(aws_account=settings.aws_account)
     settings.data_dir = settings.data_dir.format(aws_account=settings.aws_account)
+
+    if args.custom_prompts:
+        settings.custom_prompts = True
+    else:
+        settings.custom_prompts = False
     return settings
 
 
@@ -74,6 +79,11 @@ def parse_args() -> argparse.Namespace:
         required=False,
         help="Allow for a local path to a tokenizer.",
     )
+
+    parser.add_argument(
+        "--custom-prompts", default=False, help="Load custom prompts from a prompt folder"
+    )
+
     return parser.parse_args()