import os import re import uuid import time import signal import logging import threading import concurrent.futures from datetime import datetime from typing import Dict, List from tqdm import tqdm from collections import defaultdict from utils.file_io import load_json_file, save_json_file from utils.api import send_to_judge_model from utils.visualization import create_side_by_side_score_charts import statistics from core.scoring import ( parse_scores, compute_raw_score, compute_detailed_distribution, compute_model_level_stats, compute_cross_model_stats, build_landmark_calibration_config, apply_landmark_calibration, log_score_summary, confidence_interval_95 ) from core.scoring import compute_detailed_distribution, compute_detailed_distribution # etc from core.separability import compute_separability_metrics from core.stability import run_stability_test, compute_iteration_stability, compute_randomized_iteration_rank_stability_by_item from utils.stats import clamp from utils.state import should_exit, executor from utils.stats import normalize def process_sample(model_name: str, iteration_key: str, item_id: str, item_text: str, prompt_template: str, run_key: str, runs: Dict, runs_file: str, lock: threading.Lock, judge_model: str, save_raw_judge_output: bool): """Process a single sample, retrying failed or empty results.""" global should_exit if should_exit: return text_len = len(item_text) run_data = runs.get(run_key, {}) results = run_data.get("results", {}) model_dict = results.setdefault(model_name, {}) iteration_dict = model_dict.setdefault(iteration_key, {}) existing_item = iteration_dict.get(item_id, {}) if (existing_item and "aggregated_score_raw" in existing_item and existing_item.get("parsed_scores") and len(existing_item["parsed_scores"]) >= 10 and existing_item["aggregated_score_raw"] > 0.0): return try: final_prompt = prompt_template.replace("[TEST MODEL RESPONSE]", item_text) final_prompt = final_prompt.replace("[TEST MODEL RESPONSE END]", "") messages = [{"role": "user", "content": final_prompt}] judge_response = send_to_judge_model(messages, judge_model=judge_model) extracted_scores = parse_scores(judge_response) raw_score = compute_raw_score(extracted_scores) with lock: storage_dict = { "parsed_scores": extracted_scores, "timestamp": datetime.now().isoformat(), "text_length": text_len } if raw_score is not None: storage_dict["aggregated_score_raw"] = raw_score if save_raw_judge_output: storage_dict["judge_response"] = judge_response iteration_dict[item_id] = storage_dict runs[run_key]["results"][model_name][iteration_key] = iteration_dict save_json_file(runs, runs_file) if raw_score is not None: logging.debug(f"Processed {model_name}/{iteration_key}/{item_id}, raw score: {raw_score:.2f}") else: logging.warning(f"Failed to parse enough scores for {model_name}/{iteration_key}/{item_id}") except Exception as e: logging.error(f"Error processing item {model_name}/{iteration_key}/{item_id}: {str(e)}") with lock: iteration_dict[item_id] = { "error": str(e), "timestamp": datetime.now().isoformat() } if "errors" not in runs[run_key]: runs[run_key]["errors"] = [] runs[run_key]["errors"].append({ "model": model_name, "iteration": iteration_key, "item_id": item_id, "error": str(e) }) save_json_file(runs, runs_file) def finalize_scores_and_compute_judgemark(runs: dict, run_key: str, samples_data: dict): """ Compute metrics for both raw and calibrated scores, including stability tests, normalized components, and detailed distributions. Now also returns a final_judgemark_score for BOTH raw and calibrated statistics. """ run_data = runs[run_key] results = run_data.get("results", {}) # 1. Collect raw scores, compute calibration, store calibrated values raw_scores_by_model_all = defaultdict(list) raw_scores_by_model_by_iter = defaultdict(lambda: defaultdict(list)) calibrated_scores_by_model_all = defaultdict(list) calibrated_scores_by_model_by_iter = defaultdict(lambda: defaultdict(list)) lengths_by_model = {} # -- Collect raw scores for model_name, iteration_data in results.items(): if not isinstance(iteration_data, dict): continue lengths = [] for it_key, it_val in iteration_data.items(): if it_key == "__model_stats__": continue if not isinstance(it_val, dict): continue for item_id, item_info in it_val.items(): if (isinstance(item_info, dict) and "aggregated_score_raw" in item_info): raw_score = item_info["aggregated_score_raw"] # Collect raw score globally raw_scores_by_model_all[model_name].append(raw_score) # Collect raw score by iteration raw_scores_by_model_by_iter[model_name][it_key].append(raw_score) # Track text length for analyzing text = (samples_data.get(model_name, {}) .get("samples", {}) .get(it_key, {}) .get(item_id, "")) lengths.append(len(text)) if len(raw_scores_by_model_all[model_name]) > 0: lengths_by_model[model_name] = lengths # 2. Distribution + calibration all_raw_scores = [s for scores in raw_scores_by_model_all.values() for s in scores] run_data["raw_score_distribution"] = compute_detailed_distribution(all_raw_scores) calibration_config = build_landmark_calibration_config(all_raw_scores, [0, 3, 5, 7, 10]) run_data["calibration_config"] = calibration_config # Apply calibration for model_name, iteration_data in results.items(): if not isinstance(iteration_data, dict): continue # Flatten model's raw scores, calibrate them raw_list = raw_scores_by_model_all[model_name] calibrated = [apply_landmark_calibration(s, calibration_config) for s in raw_list] # Re-walk iteration_data to assign each calibration back idx = 0 for it_key, it_val in iteration_data.items(): if it_key == "__model_stats__": continue if not isinstance(it_val, dict): continue for item_id, item_info in it_val.items(): if (isinstance(item_info, dict) and "aggregated_score_raw" in item_info): item_info["aggregated_score_calibrated"] = calibrated[idx] idx += 1 # Update calibrated_scores_by_model_by_iter in the same breakdown idx2 = 0 for it_key in raw_scores_by_model_by_iter[model_name]: count_for_iter = len(raw_scores_by_model_by_iter[model_name][it_key]) these_cals = calibrated[idx2 : idx2 + count_for_iter] calibrated_scores_by_model_by_iter[model_name][it_key].extend(these_cals) idx2 += count_for_iter # Populate the single flattened list of calibrated scores calibrated_scores_by_model_all[model_name].extend(calibrated) # 3. Calibrated distributions all_calibrated_scores = [ s for scores in calibrated_scores_by_model_all.values() for s in scores ] run_data["calibrated_score_distribution"] = compute_detailed_distribution(all_calibrated_scores) # 4. Model-level stats run_data["raw_model_stats"] = compute_model_level_stats(raw_scores_by_model_all, lengths_by_model) run_data["calibrated_model_stats"] = compute_model_level_stats(calibrated_scores_by_model_all, lengths_by_model) # 5. Cross-model stats run_data["raw_cross_model_stats"] = compute_cross_model_stats( scores_by_model_all=raw_scores_by_model_all, scores_by_model_by_iter=raw_scores_by_model_by_iter ) run_data["calibrated_cross_model_stats"] = compute_cross_model_stats( scores_by_model_all=calibrated_scores_by_model_all, scores_by_model_by_iter=calibrated_scores_by_model_by_iter ) # 6. Separability metrics compute_separability_metrics(run_data, raw_scores_by_model_all, label="raw") compute_separability_metrics(run_data, calibrated_scores_by_model_all, label="calibrated") # 8. Compute iteration stability for raw & calibrated compute_iteration_stability(run_data, label="raw") compute_iteration_stability(run_data, label="calibrated") random_tau_raw = compute_randomized_iteration_rank_stability_by_item(run_data, label="raw", n_shuffles=1000) random_tau_cal = compute_randomized_iteration_rank_stability_by_item(run_data, label="calibrated", n_shuffles=1000) logging.info("Score stability (RAW)") logging.info(f"Randomized average Kendall's tau (raw): {random_tau_raw:.3f}") logging.info("Score stability (CALIBRATED)") logging.info(f"Randomized average Kendall's tau (calibrated): {random_tau_cal:.3f} " f"({run_data['calibrated_cross_model_stats']['kendall_tau']})") # 9. Compute the final Judgemark scores (one using raw stats, one using calibrated) # -- (A) RAW Judgemark # Pull out raw stats + separability metrics raw_stats = run_data["raw_cross_model_stats"] raw_norm = raw_stats["normalized_components"] # "std_dev", "kw_stat", etc. # Add your own normalization steps as needed raw_emd = run_data["separability_metrics"]["raw"]["emd"]["average"] raw_emd_norm = normalize(raw_emd, 0, 4) raw_overlap_mag = run_data["separability_metrics"]["raw"]["ci99_overlap_magnitude_sum"] raw_overlap_mag_norm = normalize(raw_overlap_mag, 0, 26, False) raw_norm["ci99_overlap_magnitude_sum_norm"] = raw_overlap_mag_norm # Range of raw model means raw_score_range = ( max(run_data["raw_model_stats"][model]["mean"] for model in run_data["raw_model_stats"]) - min(run_data["raw_model_stats"][model]["mean"] for model in run_data["raw_model_stats"]) ) run_data["raw_score_range"] = raw_score_range raw_score_range_norm = normalize(raw_score_range, 0, 8) raw_norm["raw_score_range_norm"] = raw_score_range_norm # Add Kendall's tau from the randomization-based stability measure raw_norm["kendall_tau_bootstrapped"] = normalize(random_tau_raw, 0.4, 1.0) # compute an aggregated separability metric raw_separability = ( raw_norm["std_dev"] # std deviation *between* models (separability) + raw_norm["kw_stat"] # kruskal-wallis (separability) + raw_norm["ci99_overlap_magnitude_sum_norm"] # confidence interval overlap between adjacently ranked models (separability) + raw_norm["raw_score_range_norm"] # range of assigned scores (separability) + run_data["separability_metrics"]["raw"]["modulated_ci95"] # average ci95 per model scored (score stability + separability) + raw_emd_norm # earth-movers distance (separability) ) / 6.0 # Combine into final raw Judgemark final_score_raw = ( raw_norm["kendall_tau_bootstrapped"] # correlation between iterations (ranking stability) + raw_norm["kendall_tau"] # correlation with lmsys arena score (corr to human pref) + 4 * raw_separability # aggregate of separability metrics ) / 6.0 run_data["final_judgemark_score_elements_raw"] = { "norm_stability_between_iterations": raw_norm["kendall_tau_bootstrapped"], "norm_correlation_with_lmsys_arena": raw_norm["kendall_tau"], "norm_std_dev_between_models": raw_norm["std_dev"], "norm_kruskall_wallis": raw_norm["kw_stat"], "norm_ci99_adjacent_overlap": raw_norm["ci99_overlap_magnitude_sum_norm"], "norm_score_range": raw_norm["raw_score_range_norm"], "norm_intra_model_ci95": run_data["separability_metrics"]["raw"]["modulated_ci95"], "norm_earth_movers_distance": raw_emd_norm } run_data["final_judgemark_score_raw"] = final_score_raw # -- (B) Calibrated Judgemark cal_stats = run_data["calibrated_cross_model_stats"] norm = cal_stats["normalized_components"] emd_norm = normalize(run_data["separability_metrics"]["calibrated"]["emd"]["average"], 0, 4) overlap_magnitude_norm = normalize( run_data["separability_metrics"]["calibrated"]["ci99_overlap_magnitude_sum"], 0, 26, False ) norm["ci99_overlap_magnitude_sum_norm"] = overlap_magnitude_norm # Range of calibrated model means calibrated_score_range = ( max(run_data["calibrated_model_stats"][model]["mean"] for model in run_data["calibrated_model_stats"]) - min(run_data["calibrated_model_stats"][model]["mean"] for model in run_data["calibrated_model_stats"]) ) run_data["calibrated_score_range"] = calibrated_score_range calibrated_score_range_norm = normalize(calibrated_score_range, 0, 8) norm["calibrated_score_range_norm"] = calibrated_score_range_norm # Kendall's tau from the randomized stability measure norm["kendall_tau_bootstrapped"] = normalize(random_tau_cal, 0.4, 1.0) # compute an aggregated separability metric calibrated_separability = ( norm["std_dev"] # std deviation *between* models (separability) + norm["kw_stat"] # kruskal-wallis (separability) + norm["ci99_overlap_magnitude_sum_norm"] # confidence interval overlap between adjacently ranked models (separability) + norm["calibrated_score_range_norm"] # range of assigned scores (separability) + run_data["separability_metrics"]["calibrated"]["modulated_ci95"] # average ci95 per model scored (score stability + separability) + emd_norm # earth-movers distance (separability) ) / 6.0 final_score_calibrated = ( norm["kendall_tau_bootstrapped"] # correlation between iterations (ranking stability) + norm["kendall_tau"] # correlation with lmsys arena score (corr to human pref) + 4 * calibrated_separability # aggregate of separability metrics ) / 6.0 run_data["final_judgemark_score_elements_calibrated"] = { "norm_stability_between_iterations": norm["kendall_tau_bootstrapped"], "norm_correlation_with_lmsys_arena": norm["kendall_tau"], "norm_std_dev_between_models": norm["std_dev"], "norm_kruskall_wallis": norm["kw_stat"], "norm_ci99_adjacent_overlap": norm["ci99_overlap_magnitude_sum_norm"], "norm_score_range": norm["calibrated_score_range_norm"], "norm_intra_model_ci95": run_data["separability_metrics"]["calibrated"]["modulated_ci95"], "norm_earth_movers_distance": norm } run_data["final_judgemark_score"] = final_score_calibrated # 10. Create visualizations + logs create_side_by_side_score_charts(run_data, run_data["judge_model"], samples_data) log_score_summary( "RAW SCORES", run_data["raw_cross_model_stats"], run_data["raw_model_stats"] ) log_score_summary( "CALIBRATED SCORES", run_data["calibrated_cross_model_stats"], run_data["calibrated_model_stats"] ) logging.info(f"Final Judgemark (raw) = {final_score_raw:.3f}") logging.info(f"Final Judgemark (cal) = {final_score_calibrated:.3f}") def sanitize_model_name(name: str) -> str: """Sanitize judge model name for use in the run key.""" return re.sub(r'[^a-zA-Z0-9_-]+', '_', name) def run_judgemark_v2( judge_model: str, samples_file: str, prompts_file: str, runs_file: str, num_threads: int, run_id: str = None, save_raw_judge_output: bool = False ) -> str: global executor, should_exit logging.info(f"Starting Judgemark-v2 using judge model: {judge_model}") runs = load_json_file(runs_file) # Form the run key using run_id + "__" + sanitized judge model sanitized_jm = sanitize_model_name(judge_model) base_id = run_id if run_id else str(uuid.uuid4()) run_key = f"{base_id}__{sanitized_jm}" # Load data files samples_data = load_json_file(samples_file) judge_prompts = load_json_file(prompts_file) # Initialize or get existing run data if run_key not in runs: runs[run_key] = { "judge_model": judge_model, "start_time": datetime.now().isoformat(), "status": "running", "samples_file": samples_file, "prompts_file": prompts_file, "results": {} } save_json_file(runs, runs_file) run_data = runs[run_key] items_to_process = [] # If run exists, scan for items needing retry if "results" in run_data: results = run_data.get("results", {}) # Scan all possible items for model_name, model_info in samples_data.items(): samples_dict = model_info.get("samples", {}) for iteration_key, iteration_items in samples_dict.items(): for item_id, item_text in iteration_items.items(): # Check if this item needs processing existing_result = (results.get(model_name, {}) .get(iteration_key, {}) .get(item_id, {})) needs_retry = ( not existing_result or not existing_result.get("parsed_scores") or len(existing_result.get("parsed_scores", {})) < 10 or existing_result.get("aggregated_score_raw", 0.0) == 0.0 or "error" in existing_result ) if needs_retry: items_to_process.append({ "model_name": model_name, "iteration_key": iteration_key, "item_id": item_id, "item_text": item_text, "prompt_template": judge_prompts.get(item_id, "") }) if items_to_process: logging.info(f"Found {len(items_to_process)} items to process in existing run {run_key}") else: logging.info(f"No items to process in existing run {run_key}") else: # New run - process all items for model_name, model_info in samples_data.items(): samples_dict = model_info.get("samples", {}) for iteration_key, iteration_items in samples_dict.items(): print(iteration_key) for item_id, item_text in iteration_items.items(): items_to_process.append({ "model_name": model_name, "iteration_key": iteration_key, "item_id": item_id, "item_text": item_text, "prompt_template": judge_prompts.get(item_id, "") }) # Ensure concurrency lock lock = threading.Lock() # Process any items that need retrying all_futures = [] try: with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as exec_: executor = exec_ if items_to_process: # Process all items (either retries or new run) for item in items_to_process: if should_exit: break all_futures.append( executor.submit( process_sample, item["model_name"], item["iteration_key"], item["item_id"], item["item_text"], item["prompt_template"], run_key, runs, runs_file, lock, judge_model, save_raw_judge_output ) ) # Display progress bar for tasks for f in tqdm(concurrent.futures.as_completed(all_futures), total=len(all_futures), desc="Judging", leave=True): if should_exit: break try: f.result() except Exception as exc: logging.error(f"Exception in worker thread: {exc}") except KeyboardInterrupt: logging.warning("KeyboardInterrupt caught in main thread.") should_exit = True time.sleep(0.1) finally: # Mark run as interrupted or completed status = "interrupted" if should_exit else "completed" runs[run_key]["status"] = status runs[run_key]["end_time"] = datetime.now().isoformat() if not should_exit: # Run stability test if False: run_stability_test( run_data, judge_model, judge_prompts, samples_data, runs, runs_file, lock, num_threads ) # Compute final stats finalize_scores_and_compute_judgemark(runs, run_key, samples_data) # Save final save_json_file(runs, runs_file) if executor: logging.info("Shutting down executor") executor.shutdown(wait=False) executor = None logging.info(f"Judgemark-v2 run {run_key} ended with status: {status}") return run_key