Merge pull request #15 from amazon-science/llm-judge-by-model-and-tone

llm judge when there are multiple models and tones

Author: gabriben

README.md

@@ -106,12 +106,13 @@ You can use the [CloudFormation yaml](src/cloudformation.yml) to start a Sagemak
 - [x] data
 - [x] 1. data generation -> prompt library
 - [x] 2.b. LLM -> implement this in a modular way in in format_prompt_as_xml
+- [x] merge generate_all_datasets and generate_specific_datasets
+- [x] add a model_router.py
+- [x] uv
+- [x] from main.py to setup.py
 - [ ] transfer args to settings
-- [ ] merge generate_all_datasets and generate_specific_datasets
 - [ ] batch processing for Bedrock
 - [ ] batch processing for Sagemaker endpoint
-- [x] uv
-- [x] from main.py to setup.py
 - [ ] better sagemaker inference output parsing
 - [x] add a model_router.py
 - [ ] check if model exists in settings.toml to avoid AttributeError: 'Settings' object has no attribute 'ENDPOINT_TYPE'

config/settings.toml

@@ -1,7 +1,7 @@
 [default]
 region = 'us-east-1'
-data_dir = "./data" 
-# 's3://llm-finetune-us-east-1-{aws_account}/eval/tones/'
+data_dir = 's3://llm-finetune-us-east-1-{aws_account}/eval/tones/'
+#  "./data" 
 deploy_bucket_name = 'llm-finetune-us-east-1-{aws_account}'
 deploy_bucket_prefix = 'models'
 sagemaker_execution_role_arn = 'arn:aws:iam::{aws_account}:role/sagemaker-execution-role-us-east-1'

src/wraval/actions/action_llm_judge.py

@@ -4,6 +4,7 @@
 #
 import pandas as pd
 from typing import List, Dict, Optional
+from itertools import product
 from dynaconf import Dynaconf
 from .data_utils import write_dataset, load_latest_dataset
 from .prompts_judge import generate_input_prompt, generate_system_prompt, get_rubric, rewrite_prompt
@@ -24,7 +25,7 @@ def extract_score(text: str) -> Optional[int]:
     match = re.search(r"<score>(\d+)</score>", text)
     return int(match.group(1)) if match else None
 
-def validate_dataset(d: pd.DataFrame) -> bool:
+def validate_dataset(results: pd.DataFrame) -> bool:
     """Validate required columns exist in dataset.
     
     Args:
@@ -34,14 +35,14 @@ def validate_dataset(d: pd.DataFrame) -> bool:
         True if valid, False otherwise
     """
     required_columns = {"synthetic_data", "rewrite", "tone"}
-    if not all(col in d.columns for col in required_columns):
+    if not all(col in results.columns for col in required_columns):
         print(f"Missing required columns. Required: {required_columns}")
         return False
     return True
 
 def process_tone_data(
     settings: Dynaconf, 
-    d: pd.DataFrame,
+    results: pd.DataFrame,
     tone: str,
     model_name: str,
     client: boto3.client,
@@ -63,14 +64,14 @@ def process_tone_data(
     if settings.custom_prompts == True:
         from wraval.custom_prompts.prompts_judge import generate_input_prompt, generate_system_prompt
 
-    dmt = d[d.tone == tone].copy()
+    temp_results = results.copy()
     rubrics = list(tone_rubrics.keys())
 
     # Generate prompts
     user_prompts = []
     sys_prompts = []
 
-    for q, a in zip(dmt["synthetic_data"], dmt["rewrite"]):
+    for q, a in zip(temp_results["synthetic_data"], temp_results["rewrite"]):
         for rubric in rubrics:
             user_prompts.append(generate_input_prompt(q, a, tone))
             sys_prompts.append(generate_system_prompt(tone_rubrics[rubric]))
@@ -88,16 +89,16 @@ def process_tone_data(
 
     # Process scores
     for i, rubric in enumerate(rubrics):
-        dmt[rubric] = completions[i::len(rubrics)]
-        dmt[f'{rubric}_score'] = dmt[rubric].apply(extract_score)
+        temp_results[rubric] = completions[i::len(rubrics)]
+        temp_results[f'{rubric}_score'] = temp_results[rubric].apply(extract_score)
 
     # Move all score columns to the right
     score_columns = [f'{r}_score' for r in rubrics]
-    other_columns = [col for col in dmt.columns if col not in score_columns]
-    dmt = dmt[other_columns + score_columns]
+    other_columns = [col for col in temp_results.columns if col not in score_columns]
+    temp_results = temp_results[other_columns + score_columns]
 
-    dmt['overall_score'] = dmt[score_columns].mean(axis=1)
-    return dmt
+    temp_results['overall_score'] = temp_results[score_columns].mean(axis=1)
+    return temp_results
 
 def judge(
     settings: Dynaconf,
@@ -121,30 +122,39 @@ def judge(
         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")
+        results = load_latest_dataset(settings.data_dir)
+        print(f"Loaded dataset with {len(results)} rows")
     except FileNotFoundError:
         print("No dataset found. Please generate data first.")
         return
 
-    if not validate_dataset(d):
+    if not validate_dataset(results):
         return
 
-    tones = d["tone"].unique()
+    tones = results["tone"].unique()
+    inf_models = results["inference_model"].unique()
     print(f"Found tones: {tones}")
+    print(f"Found inference_models: {inf_models}")
+
+    if settings.type != "all":
+        tones = [settings.type]
 
-    for tone in tones:
-        print(f"\n{'='*20}\n{tone}\n{'='*20}")
+    # Process each tone-model combination that needs scoring
+    for tone, inf_model in product(tones, inf_models):
+        mask = (results.inference_model == inf_model) & (results.tone == tone)
+        # check if any score is missing for this inference model and this tone
+        # If yes, run the eval below
+        if not results[mask].overall_score.isna().any():
+            continue
+            
+        print(f"\n{'='*20}\n{tone} tone\nfor inference model {inf_model}\n{'='*20}")
 
         tone_rubrics = get_rubric(tone.upper())
-        dmt = process_tone_data(settings, d, tone, model_name, client, tone_rubrics)
-        
-        # Update main dataframe
-        mask = (d.tone == tone)
-        d.loc[mask, dmt.columns] = dmt.values
+        temp_results = process_tone_data(settings, results[mask], tone, model_name, client, tone_rubrics)
+        results.loc[mask, temp_results.columns] = temp_results.values
 
     # Save results
-    write_dataset(d, settings.data_dir, "all-tones", "csv")
+    write_dataset(results, settings.data_dir, "all", "csv")
 
 def rewrite_judge(
     model_id: str,
@@ -163,12 +173,12 @@ def rewrite_judge(
     Returns:
         DataFrame with input, output, and scores
     """
-    d = pd.DataFrame({'input': queries, 'output': answers})
+    results = pd.DataFrame({'input': queries, 'output': answers})
     prompts = [rewrite_prompt(q, a) for q, a in zip(queries, answers)]
-    d['rewrite_score'] = batch_get_bedrock_completions(
+    results['rewrite_score'] = batch_get_bedrock_completions(
         model_id, 
         bedrock_client,
         prompts,
         max_concurrent=len(prompts)
     )
-    return d
+    return results

src/wraval/actions/action_results.py

@@ -48,7 +48,7 @@ def show_results(settings: Dynaconf, tone: Optional[str] = None) -> None:
         print("=" * 50)
 
         # Group by model, inference model, and tone, calculate mean of overall_score
-        grouped = d.groupby(['model', 'tone'])['overall_score'].mean()
+        grouped = d.groupby(['inference_model', 'tone'])['overall_score'].mean()
 
         # Normalize scores to 0-100 scale
         normalized = normalize_scores(grouped)

src/wraval/actions/data_utils.py

@@ -1,5 +1,5 @@
 import os
-from datetime import datetime
+from datetime import datetime, UTC
 import pandas as pd
 import boto3
 import tempfile
@@ -45,7 +45,7 @@ def write_dataset_local(
 
 
 def add_timestamp_to_file_prefix(file_prefix, format):
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
     return f"{file_prefix}-{timestamp}.{format.lower()}"
 
 

src/wraval/actions/model_router.py

@@ -44,14 +44,15 @@ def get_completion(self, queries: List[str]) -> List[str]:
 class SageMakerRouter(HuggingFaceModelRouter):
     def __init__(self, master_sys_prompt, settings):
         super().__init__(master_sys_prompt, settings)
+        self.model_name = settings.model
 
     def get_completion(self, queries: List[str]) -> List[str]:
         prompts = [
             format_prompt(text, self.master_sys_prompt, self.tokenizer, type="hf")
             for text in queries
         ]
         return [
-            invoke_sagemaker_endpoint({"inputs": prompt}) for prompt in tqdm(prompts)
+            invoke_sagemaker_endpoint({"inputs": prompt}, self.model_name) for prompt in tqdm(prompts)
         ]