llm judge

Author: gabriben

.gitignore

@@ -7,3 +7,4 @@ src/wraval.egg-info/
 **__pycache__/
 prompts/*
 .idea
+src/wraval/custom_prompts/*

config/settings.toml

@@ -1,6 +1,7 @@
 [default]
 region = 'us-east-1'
-data_dir = 's3://llm-finetune-us-east-1-{aws_account}/eval/tones/'
+data_dir = "./data"
+# 's3://llm-finetune-us-east-1-{aws_account}/eval/tones/'
 
 [haiku-3]
 model = 'anthropic.claude-3-haiku-20240307-v1:0'

src/wraval/actions/action_llm_judge.py

@@ -39,6 +39,7 @@ def validate_dataset(d: pd.DataFrame) -> bool:
     return True
 
 def process_tone_data(
+    settings: Dynaconf, 
     d: pd.DataFrame,
     tone: str,
     model_name: str,
@@ -61,22 +62,21 @@ def process_tone_data(
     rubrics = list(tone_rubrics.keys())
 
     # Generate prompts
-    prompts = []
+    user_prompts = []
+    sys_prompts = []
+
     for q, a in zip(dmt["synthetic_data"], dmt["rewrite"]):
         for rubric in rubrics:
-            prompts.append((
-                generate_system_prompt(tone_rubrics[rubric]),
-                generate_input_prompt(q, a, tone)
-            ))
+            user_prompts.append(generate_input_prompt(q, a, tone))
+            sys_prompts.append(generate_system_prompt(tone_rubrics[rubric]))
 
     # Get completions
-    sys_prompts, user_prompts = zip(*prompts)
+    # import pdb
+    # pdb.set_trace()
     completions = batch_get_bedrock_completions(
-        model_name, 
-        client, 
+        settings,
         user_prompts, 
-        sys_prompts, 
-        max_concurrent=len(user_prompts)
+        sys_prompts
     )
 
     rubrics = [r.lower() for r in rubrics]
@@ -99,7 +99,6 @@ def judge(
     client: boto3.client,
     model_name: str,
     upload_s3: bool,
-    data_dir: str,
     endpoint_type: str = "bedrock"
 ) -> None:
     """Judge rewrites using specified model and rubrics.
@@ -113,7 +112,7 @@ def judge(
         endpoint_type: Type of endpoint to use
     """
     try:
-        d = load_latest_dataset(data_dir)
+        d = load_latest_dataset(settings.data_dir)
         print(f"Loaded dataset with {len(d)} rows")
     except FileNotFoundError:
         print("No dataset found. Please generate data first.")
@@ -129,7 +128,7 @@ def judge(
         print(f"\n{'='*20}\n{tone}\n{'='*20}")
 
         tone_rubrics = get_rubric(tone.upper())
-        dmt = process_tone_data(d, tone, model_name, client, tone_rubrics)
+        dmt = process_tone_data(settings, d, tone, model_name, client, tone_rubrics)
 
         # Update main dataframe
         mask = (d.tone == tone)

src/wraval/actions/completion.py

@@ -52,8 +52,8 @@ def get_bedrock_completion(settings, prompt, system_prompt=None):
                 )
 
             if isinstance(system_prompt, str) and len(system_prompt) > 0:
-                # converse_api_params.update({"system": [{"text": system_prompt}]})
-                converse_api_params["messages"] = [{"role": "assistant", "content": [{"text": system_prompt}]}] + converse_api_params["messages"]
+                converse_api_params.update({"system": [{"text": system_prompt}]})
+                # converse_api_params["messages"] = [{"role": "assistant", "content": [{"text": system_prompt}]}] + converse_api_params["messages"]
 
             response = bedrock_client.converse(**converse_api_params)
             return response['output']['message']['content'][0]['text']

src/wraval/actions/data_utils.py

@@ -97,6 +97,7 @@ def load_latest_dataset(data_dir: str) -> pd.DataFrame:
         bucket, prefix = parse_s3_path(data_dir)
         return load_latest_dataset_from_s3(bucket, prefix)
     else:
+
         # Local file handling
         data_dir = os.path.expanduser(data_dir)
 

src/wraval/main.py

@@ -86,7 +86,7 @@ def handle_inference(args, settings):
 
 
 def handle_judge(args, settings):
-    if args.endpoint_type == "bedrock":
+    if settings.endpoint_type == "bedrock":
         judge_model = settings.model
         client = boto3.client(
             service_name="bedrock-runtime", region_name=settings.region
@@ -100,8 +100,7 @@ def handle_judge(args, settings):
         client,
         judge_model,
         args.upload_s3,
-        settings.data_dir,
-        args.endpoint_type,
+        settings.endpoint_type,
     )