import argparse import hashlib import json from pathlib import Path import polars as pl _DESCRIPTION = """Compare and find differences to benchmark outputs. Supported inputs: * The samples jsonl from `lm_eval --log_samples --output_path FOLDER_NAME` * The output from `gsm8k/bench_sglang.py --raw-result-file FILE_NAME` (or mmlu) """ def main(args): if args.data_type == "simple_evals": df_input = _compute_df_input_mode_simple_evals(args) else: df_input = _transform_df_input(_compute_df_raw(args)) assert all( c in df_input.columns for c in ["category", "trial_index", "prompt_id", "prompt", "output", "correct"] ) df_meta = _compute_df_meta(df_input) df_correctness_per_trial = df_input.group_by( "category", "trial_index", maintain_order=True ).agg(pl.col("correct").mean()) df_correctness_delta = ( df_meta.group_by("correctness_delta").len().sort("correctness_delta") ) df_good_to_bad = df_meta.filter(pl.col("correctness_delta") < 0) df_bad_to_good = df_meta.filter(pl.col("correctness_delta") > 0) print(f"Dump output to {args.output_path}") Path(args.output_path).write_text( json.dumps( dict( df_meta=df_meta.to_dicts(), df_good_to_bad=df_good_to_bad.to_dicts(), df_bad_to_good=df_bad_to_good.to_dicts(), ), indent=4, ), ) if not args.disable_print_details: with pl.Config( fmt_str_lengths=10000, tbl_cols=-1, tbl_rows=-1, tbl_width_chars=-1, tbl_formatting="UTF8_FULL", ): print("====== Correctness per trial ======") print(df_correctness_per_trial) print( "====== Correctness Delta (-1.0 means all-right becomes all-wrong) ======" ) print(df_correctness_delta) for name, df in [ ("Good->Bad", df_good_to_bad), ("Bad->Good", df_bad_to_good), ]: print(f"====== Concrete Examples: {name} ======") print(df) def _compute_df_input_mode_simple_evals(args): return pl.concat( [ _compute_df_input_one_mode_simple_evals(**info) for info in _get_file_infos(args=args) ] ) def _compute_df_input_one_mode_simple_evals(path, category, trial_index): data = json.loads(Path(path).read_text()) rows = [] for single_eval_result in data["metadata"]["single_eval_results"]: prompt = single_eval_result["example_level_metadata"][ "actual_queried_prompt_messages" ] score = single_eval_result["score"] assert score in {0.0, 1.0}, f"{score=}" row = dict( category=category, trial_index=trial_index, prompt_id=_compute_id_from_object(prompt), prompt=json.dumps(prompt), output=single_eval_result["example_level_metadata"]["response_text"], correct=score == 1.0, ) rows.append(row) return pl.DataFrame(rows) def _compute_id_from_object(obj): if isinstance(obj, pl.Series): obj = obj.to_list() json_str = json.dumps(obj, sort_keys=True, ensure_ascii=False) return hashlib.sha256(json_str.encode("utf-8")).hexdigest() def _compute_df_raw(args): return pl.concat( [ _read_df_raw( path=info["path"], category=info["category"], trial_index=info["trial_index"], ) for info in _get_file_infos(args=args) ] ) def _get_file_infos(args): return [ dict(path=path, category=category, trial_index=trial_index) for category, paths in [ ("baseline", args.baseline_path), ("target", args.target_path), ] for trial_index, path in enumerate(paths) ] def _read_df_raw(path: str, category: str, trial_index: int): return pl.read_ndjson(path).with_columns( category=pl.lit(category), trial_index=trial_index ) def _transform_df_input(df: pl.DataFrame): if "doc_id" in df.columns: print("Transform mode: lm_eval") filter_names = df["filter"].unique(maintain_order=True).to_list() if len(filter_names) > 1: filter_name = filter_names[0] print(f"Choose {filter_name=} among {filter_names}") df = df.filter(pl.col("filter") == filter_name) df = df.select( pl.col("category"), pl.col("trial_index"), prompt_id=pl.col("doc_id"), prompt=pl.col("arguments").struct.field("gen_args_0").struct.field("arg_0"), output=pl.col("resps").list.get(0).list.get(0), correct=pl.col("exact_match").cast(bool), ) return df elif "prompt_id" in df.columns: print("Transform mode: SGLang bench") return df else: raise Exception( f"Unknown data: {df.columns}. You may need to set `--data-type` if using e.g. simple_evals." ) def _compute_df_meta(df_input: pl.DataFrame): df_input = df_input.sort("prompt_id", "category", "trial_index") df_meta = pl.DataFrame( [ _handle_one_prompt(df_one_prompt) for df_one_prompt in df_input.partition_by("prompt_id", maintain_order=True) ] ) df_meta = df_meta.with_columns( correctness_delta=pl.col("correctness_target") - pl.col("correctness_baseline"), ) df_meta = df_meta.sort("correctness_delta", "output_same_prefix_len") return df_meta def _handle_one_prompt(df_one_prompt: pl.DataFrame): assert ( len(set(_compute_id_from_object(obj) for obj in df_one_prompt["prompt"])) == 1 ) df_baseline = df_one_prompt.filter(pl.col("category") == "baseline") df_target = df_one_prompt.filter(pl.col("category") == "target") outputs_baseline = df_baseline["output"].to_list() outputs_target = df_target["output"].to_list() output_same_prefix_len = max( _compute_str_prefix_len(output_baseline, output_target) for output_baseline in outputs_baseline for output_target in outputs_target ) return dict( prompt_id=df_one_prompt[0, "prompt_id"], correctness_baseline=df_baseline["correct"].mean(), correctness_target=df_target["correct"].mean(), output_same_prefix_len=output_same_prefix_len, prompt=df_one_prompt[0, "prompt"], outputs_baseline=outputs_baseline, outputs_target=outputs_target, ) def _compute_str_prefix_len(a: str, b: str) -> int: min_len = min(len(a), len(b)) for i in range(min_len): if a[i] != b[i]: return i return min_len if __name__ == "__main__": parser = argparse.ArgumentParser(description=_DESCRIPTION) parser.add_argument("--data-type", type=str, default="auto") parser.add_argument("--baseline-path", type=str, nargs="+") parser.add_argument("--target-path", type=str, nargs="+") parser.add_argument( "--output-path", type=str, default="/tmp/text_comparator_output.json" ) parser.add_argument("--disable-print-details", action="store_true") args = parser.parse_args() main(args)