diff --git a/.gitignore b/.gitignore index 03042d4..3630a70 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ openai_batches/ models/ results/ datasets/ +llm-comparator/ /data/ __pycache__/ *egg-info/ diff --git a/README.md b/README.md index 2a9b45c..ad8ef7a 100644 --- a/README.md +++ b/README.md @@ -92,10 +92,12 @@ uv run data/download_fineweb_edu.py # 1 and 2 can be run in parallel (depends on step 0) # 1. generate QA data (3 QAs for each context) +# 000_0000[0-1] to 000_0000[8-9] for small split vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python generate_fw_edu_qa_vllm.py "00*_*" 3 # 2. Augmenting pre-training data w/ paraphrasing -uv run python data/generate_fw_edu_augment_vllm.py +# (000_00000 to 000_00009 for small split) +vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python data/generate_fw_edu_augment_vllm.py 000_00000 # === # 3 and 4 can be run in parallel (depends on step 0 and 1) @@ -109,6 +111,7 @@ uv run python data/generate_pretrain_from_fw_qa.py # the responses are not the same as the ones in the pretraining though # Example commands using gemma-2-2b-it # self-gen data for fw_qa +# 000_0000[0-1] to 000_0000[8-9] for small split uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_3/00*_*' # self-gen data for other ds listed in qa_no_fw.yaml uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_no_fw.yaml @@ -141,16 +144,16 @@ run python intx_sft.py configs/...yaml ... --from_pretrained_checkpoint=train_ou LongBench ```bash # generative -run python run_eval.py --checkpoint_path train_outputs/runs/.../pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test +WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/.../pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test # hypernet checkpoint -uv run python run_eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test +WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test # base model -uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --eval_batch_size 2 +WANDB_MODE=disabled uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --eval_batch_size 2 # base model w/o context -run uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --remove_context +WANDB_MODE=disabled uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --remove_context # # benchmark # cd LongBench/LongBench @@ -158,4 +161,30 @@ run uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datase # run python pred_ctx_to_lora.py --checkpoint_path ../../train_outputs/runs/Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782/pytorch_model.bin # run python eval_ctx_to_lora.py --model_name Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782 --checkpoint_path ../../train_outputs/runs/Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782/pytorch_model.bin +``` + +### LLM-comparator +```bash +# install nvm +curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash +nvm install 16 +nvm use 16 + +git clone https://github.com/PAIR-code/llm-comparator.git +cd llm-comparator +npm install +npm run build +# running llm-comparator webui +npm run serve + +# in another terminal +# run http server for file fetching + +# cd back to root folder first +# taken from https://stackoverflow.com/a/79135787 +alias srv='echo -e "from sys import argv as a\nfrom http.server import HTTPServer as H, SimpleHTTPRequestHandler as HH, test as t\nclass C(HH):\n def end_headers (self):\n self.send_header(a[2],a[3])\n HH.end_headers(self)\nt(C,H,port=int(a[1]))" | /usr/bin/env python3 -- - 8001 "Access-Control-Allow-Origin" "*"' +srv + +# copy-paste the relative path from the root +# e.g., http://localhost:8001/train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/eval-results-105000/comparator.json ``` \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 77fd880..79aafd0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ exclude = [ "EditingLlama", "icae_v2", "lm-evaluation-harness", + "llm-comparator", "LongBench", "scripts", "train_outputs", diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index ff3e7db..e723e2e 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -2,7 +2,7 @@ import hashlib import json import logging import random -from collections.abc import Callable, Iterator +from collections.abc import Callable from glob import glob from os import path from typing import Any @@ -40,7 +40,7 @@ def load_answers(ds_name, split): ds_kwargs = get_ds_kwargs(ds_name, split) ds = load_dataset(**ds_kwargs, trust_remote_code=True) - ds = ds.map(extract_ans, num_proc=8) + ds = ds.map(extract_ans, num_proc=8, remove_columns=ds.column_names) return ds @@ -460,7 +460,7 @@ def load_and_process_dataset( cols_to_remove = [ col for col in ds.column_names - if col not in ["context", "prompt", "response", "qas"] + if col not in ["context", "prompt", "response", "qas", "variation"] ] is_eval = split != "train" ds = ds.map( @@ -955,91 +955,91 @@ def get_assistant_start_end_indices( return assistant_ranges -def get_masked_labels( - conversation_ids: dict[str, list[Any]], - assistant_ranges: list[tuple[int, int]], -) -> Iterator[int]: - """ - Generate masked labels for conversation, masking non-assistant tokens. - NOTE: This will also includes extra tokens between assistant and user messages - in multi-turn conversations. - E.g., {assistant_msg} <|eot_id|><|start_header_id|>user<|end_header_id|> {user_msg} - for Llama 3 models. +# def get_masked_labels( +# conversation_ids: dict[str, list[Any]], +# assistant_ranges: list[tuple[int, int]], +# ) -> Iterator[int]: +# """ +# Generate masked labels for conversation, masking non-assistant tokens. +# NOTE: This will also includes extra tokens between assistant and user messages +# in multi-turn conversations. +# E.g., {assistant_msg} <|eot_id|><|start_header_id|>user<|end_header_id|> {user_msg} +# for Llama 3 models. - Args: - conversation_ids: Dictionary with tokenized conversation info - assistant_ranges: List of (start, end) indices for assistant messages +# Args: +# conversation_ids: Dictionary with tokenized conversation info +# assistant_ranges: List of (start, end) indices for assistant messages - Yields: - Token ID or IGNORE_INDEX for each position - """ - for id_, (id_s, id_e) in list( - zip(conversation_ids["input_ids"], conversation_ids["offset_mapping"]) - ): - if any(id_s >= s and id_e <= e for s, e in assistant_ranges): - yield id_ - else: - yield IGNORE_INDEX +# Yields: +# Token ID or IGNORE_INDEX for each position +# """ +# for id_, (id_s, id_e) in list( +# zip(conversation_ids["input_ids"], conversation_ids["offset_mapping"]) +# ): +# if any(id_s >= s and id_e <= e for s, e in assistant_ranges): +# yield id_ +# else: +# yield IGNORE_INDEX -def tokenize_chat_messages( - example: dict[str, Any], - tokenizer: PreTrainedTokenizerBase, - mask_assistant_inputs: bool = True, - for_kl_loss: bool = False, - tokenizer_kwargs: dict[str, Any] | None = None, -) -> dict[str, list[int]]: - """ - Tokenize chat messages and optionally mask non-assistant tokens. +# def tokenize_chat_messages( +# example: dict[str, Any], +# tokenizer: PreTrainedTokenizerBase, +# mask_assistant_inputs: bool = True, +# for_kl_loss: bool = False, +# tokenizer_kwargs: dict[str, Any] | None = None, +# ) -> dict[str, list[int]]: +# """ +# Tokenize chat messages and optionally mask non-assistant tokens. - Args: - example: Dictionary containing 'chat' and 'messages' keys - tokenizer: Tokenizer to use - mask_assistant_inputs: Whether to mask non-assistant tokens - for_kl_loss: change the column names to "chat_ids" and "chat_attn_mask" - tokenizer_kwargs: Additional arguments to pass to tokenizer +# Args: +# example: Dictionary containing 'chat' and 'messages' keys +# tokenizer: Tokenizer to use +# mask_assistant_inputs: Whether to mask non-assistant tokens +# for_kl_loss: change the column names to "chat_ids" and "chat_attn_mask" +# tokenizer_kwargs: Additional arguments to pass to tokenizer - Returns: - Dictionary containing tokenized inputs and labels - """ - # should be used only with chat models - text = example["chat"] - messages = example["messages"] - n_response = len([m for m in messages if m["role"] == "assistant"]) - if n_response != 1: - raise ValueError(f"Expected 1 assistant response. Got {n_response}.") - conversation_ids = tokenizer( - text, - # return_offsets_mapping=mask_assistant_inputs, - add_special_tokens=False, - truncation=False, - **(tokenizer_kwargs or {}), - ) - # if tokenizer_kwargs and ( - # len(conversation_ids["input_ids"]) >= tokenizer_kwargs["max_length"] - # ): - # raise ValueError( - # f"Conversation length {len(conversation_ids['input_ids'])} exceeds max length {tokenizer_kwargs['max_length']}" - # ) +# Returns: +# Dictionary containing tokenized inputs and labels +# """ +# # should be used only with chat models +# text = example["chat"] +# messages = example["messages"] +# n_response = len([m for m in messages if m["role"] == "assistant"]) +# if n_response != 1: +# raise ValueError(f"Expected 1 assistant response. Got {n_response}.") +# conversation_ids = tokenizer( +# text, +# # return_offsets_mapping=mask_assistant_inputs, +# add_special_tokens=False, +# truncation=False, +# **(tokenizer_kwargs or {}), +# ) +# # if tokenizer_kwargs and ( +# # len(conversation_ids["input_ids"]) >= tokenizer_kwargs["max_length"] +# # ): +# # raise ValueError( +# # f"Conversation length {len(conversation_ids['input_ids'])} exceeds max length {tokenizer_kwargs['max_length']}" +# # ) - if mask_assistant_inputs: - assistant_ranges = get_assistant_start_end_indices( - messages, tokenizer, conversation_ids["input_ids"] - ) - # labels = get_masked_labels(conversation_ids, assistant_ranges) - labels = [ - id_ if any(s <= i < e for s, e in assistant_ranges) else IGNORE_INDEX - for i, id_ in enumerate(conversation_ids["input_ids"]) - ] - conversation_ids["labels"] = labels - # del conversation_ids["offset_mapping"] - else: - conversation_ids["labels"] = conversation_ids["input_ids"] - if for_kl_loss: - conversation_ids["chat_ids"] = conversation_ids.pop("input_ids") - conversation_ids["chat_attn_mask"] = conversation_ids.pop("attention_mask") - conversation_ids["chat_labels"] = conversation_ids.pop("labels") - return conversation_ids +# if mask_assistant_inputs: +# assistant_ranges = get_assistant_start_end_indices( +# messages, tokenizer, conversation_ids["input_ids"] +# ) +# # labels = get_masked_labels(conversation_ids, assistant_ranges) +# labels = [ +# id_ if any(s <= i < e for s, e in assistant_ranges) else IGNORE_INDEX +# for i, id_ in enumerate(conversation_ids["input_ids"]) +# ] +# conversation_ids["labels"] = labels +# # del conversation_ids["offset_mapping"] +# else: +# conversation_ids["labels"] = conversation_ids["input_ids"] +# if for_kl_loss: +# conversation_ids["chat_ids"] = conversation_ids.pop("input_ids") +# conversation_ids["chat_attn_mask"] = conversation_ids.pop("attention_mask") +# conversation_ids["chat_labels"] = conversation_ids.pop("labels") +# return conversation_ids def tokenize_pretrain( @@ -1180,7 +1180,7 @@ def pack( "max_packed_size": max_packed_size, }, batched=True, - batch_size=1_000_000, + batch_size=250_000, num_proc=num_proc, remove_columns=ds.column_names, ) diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py index 3267bf8..3e31272 100644 --- a/src/ctx_to_lora/eval_utils.py +++ b/src/ctx_to_lora/eval_utils.py @@ -107,7 +107,7 @@ def compute_qa_f1_score( label_words = normalized_label.split() score = max(score, f1_score(prediction_words, label_words)) res.append(score) - return dict(qa_f1=np.mean(res)) + return dict(qa_f1_score=np.mean(res)), dict(qa_f1_score=res) def add_longbench_tasks(ds_names: list[str]) -> None: @@ -120,7 +120,12 @@ def add_longbench_tasks(ds_names: list[str]) -> None: ds_names += LONGBENCH_E_TASKS -def save_generated_text(samples: list[dict], output_dir: str, split: str) -> None: +def save_generated_text( + samples: list[dict], + per_sample_metric: dict[str, list[float]], + output_dir: str, + split: str, +) -> None: """Save generated text samples to JSONL file.""" os.makedirs(output_dir, exist_ok=True) # Create any necessary subdirectories if split contains path separators @@ -128,8 +133,12 @@ def save_generated_text(samples: list[dict], output_dir: str, split: str) -> Non split_dir = os.path.join(output_dir, os.path.dirname(split)) os.makedirs(split_dir, exist_ok=True) + metric_keys = list(per_sample_metric.keys()) + assert len(metric_keys) == 1 + metric_name = metric_keys[0] with open(f"{output_dir}/{split}_generated_text.jsonl", "w") as f: - for sample in samples: + for sample, metric_val in zip(samples, per_sample_metric[metric_name]): + sample[f"{metric_name}"] = metric_val f.write(json.dumps(sample) + "\n") @@ -339,10 +348,15 @@ def create_metrics_csv( # Create DataFrame if rows: - df = pd.DataFrame(rows) + new_df = pd.DataFrame(rows) + # define categories so that they're sorted properly + new_df["group_len"] = pd.Categorical( + new_df["group_len"], + categories=[f"{i}-{j}" for (i, j) in LENGTH_BINS] + ["overall"], + ) # Sort by tasks, then by length group - df = df.sort_values(["tasks"]).reset_index(drop=True) + new_df = new_df.sort_values(["tasks", "group_len"]).reset_index(drop=True) # Construct filename csv_filename = "evaluation_results" @@ -354,7 +368,33 @@ def create_metrics_csv( csv_path = os.path.join(output_dir, csv_filename) os.makedirs(os.path.dirname(csv_path), exist_ok=True) + + # Check if CSV already exists and merge if it does + if os.path.exists(csv_path): + try: + existing_df = pd.read_csv(csv_path) + + # Remove existing rows with the same model_name, group_len, and tasks + # to avoid duplicates when updating + mask = ~existing_df["tasks"].isin(new_df["tasks"]) + existing_df = existing_df[mask] + + # Concatenate existing and new data + df = pd.concat([existing_df, new_df], ignore_index=True) + + # Sort by tasks, then by length group + df = df.sort_values(["tasks"]).reset_index(drop=True) + + print(f"Updated existing CSV with {len(new_df)} new rows") + except Exception as e: + print(f"Warning: Could not read existing CSV ({e}), creating new file") + df = new_df + else: + df = new_df + print(f"Created new CSV with {len(new_df)} rows") + df.to_csv(csv_path, index=False) + print(f"Evaluation results saved to: {csv_path}") else: print("No evaluation data found to save to CSV") @@ -449,12 +489,14 @@ def eval_generation( if ds_name in CLOSED_QA_DATASETS: print("Computing QA F1 Score") - qa_f1_metric = compute_qa_f1_score(pred_texts, answers_list) + qa_f1_metric, per_sample_metric = compute_qa_f1_score( + pred_texts, answers_list + ) for k, v in qa_f1_metric.items(): eval_result.metrics[f"{split_name}_{k}"] = v eval_result.metrics[f"{split_name}_num_samples_{k}"] = n else: - rouge_metrics = compute_rouge(pred_texts, label_texts) + rouge_metrics, per_sample_metric = compute_rouge(pred_texts, label_texts) for k, v in rouge_metrics.items(): eval_result.metrics[f"{split_name}_{k}"] = v eval_result.metrics[f"{split_name}_num_samples_{k}"] = n @@ -500,7 +542,7 @@ def eval_generation( group_answers.append([txt["label"]]) break - group_qa_f1_metric = compute_qa_f1_score( + group_qa_f1_metric, _ = compute_qa_f1_score( data["generated"], group_answers ) for k, v in group_qa_f1_metric.items(): @@ -509,7 +551,7 @@ def eval_generation( f"{split_name}_num_samples_{k}_len_{group_key}" ] = data["count"] else: - group_rouge_metrics = compute_rouge( + group_rouge_metrics, _ = compute_rouge( data["generated"], data["label"] ) for k, v in group_rouge_metrics.items(): @@ -521,6 +563,7 @@ def eval_generation( save_generated_text( decoded_txts, + per_sample_metric, split=split_name, output_dir=eval_trainer.args.output_dir, ) @@ -692,7 +735,7 @@ def evaluate( for ds_name, ds in datasets.items(): val_indices = np.random.permutation(len(ds))[:max_eval_samples_per_ds] datasets[ds_name] = ds.select(val_indices) - answers[ds_name] = answers.select(val_indices) + answers[ds_name] = answers[ds_name].select(val_indices) print(f"Datasets: {datasets}") print(f"Answers: {answers}") @@ -827,7 +870,6 @@ def run_eval( test_ds_names=[], remove_context=remove_context, ) - print(args) setup_logging(args.logging_dir) # Override dataset names if provided via CLI diff --git a/src/ctx_to_lora/metrics.py b/src/ctx_to_lora/metrics.py index 56c0409..974f6bc 100644 --- a/src/ctx_to_lora/metrics.py +++ b/src/ctx_to_lora/metrics.py @@ -32,9 +32,10 @@ def compute_rouge(pred_texts, label_texts): scores = scorer.score(pred_text, label_text) for k, v in scores.items(): out[f"{k}.f1"].append(v.fmeasure) + out_mean = dict() for k in out: - out[k] = np.mean(out[k]) - return out + out_mean[k] = np.mean(out[k]) + return out_mean, out @torch.inference_mode() @@ -85,8 +86,8 @@ class Evaluator: self.reset() def reset(self): - self.accum_metrics = defaultdict(list) - self.count = defaultdict(list) + self.accum_metrics = defaultdict(lambda: list((0,))) + self.count = defaultdict(lambda: list((0,))) def update(self, shift_logits, shift_labels, valid_masks, lengths=None): for metric_fn in self.metric_fns: @@ -137,7 +138,7 @@ class Evaluator: def compute(self): # Get result across entire eval set result = { - k: np.sum(v) / np.sum(self.count[k]) if np.sum(v) != 0 else "None" + k: np.sum(v) / np.sum(self.count[k]) if len(v) > 1 else "None" for k, v in self.accum_metrics.items() } # for k, v in self.count.items(): diff --git a/src/ctx_to_lora/viz_utils.py b/src/ctx_to_lora/viz_utils.py new file mode 100644 index 0000000..df17ff1 --- /dev/null +++ b/src/ctx_to_lora/viz_utils.py @@ -0,0 +1,191 @@ +from glob import glob + +import pandas as pd + +# convert to llm-comparator format +"""{ + "metadata": { + "source_path": "Any string for your records (e.g., run id)", + "custom_fields_schema": [] + }, + "models": [ + {"name": "Short name of your first model"}, + {"name": "Short name of your second model"} + ], + "examples": [ + { + "input_text": "This is a prompt.", + "tags": ["Math"], # A list of keywords for categorizing prompts + "output_text_a": "Response to the prompt from the first model (A)", + "output_text_b": "Response to the prompt from the other model (B)", + "score": -1.25, # Score from the judge LLM + "individual_rater_scores": [], + "custom_fields": {} + }, + { + "input_text": "This is a next prompt.", + ... + } + ] +} +""" + + +def convert_to_llm_comparator_format( + eval_checkpoint_dir: str, + model_dir: str, + compare_to_no_context: bool = False, +) -> dict: + # return { + # "metadata": {"source_path": source_path, "custom_fields_schema": []}, + # "models": [{"name": model_a_name}, {"name": model_b_name}], + # "examples": examples, + # } + + checkpoint_name = "/".join(eval_checkpoint_dir.strip("/").split("/")[-2:]) + files_checkpoint = sorted( + glob(f"{eval_checkpoint_dir}/**/*generated*.jsonl", recursive=True) + ) + files_checkpoint = {file.split("/")[-1]: file for file in files_checkpoint} + + eval_model_dir = f"eval_results/{model_dir}" + if compare_to_no_context: + files_model = sorted( + glob(f"{eval_model_dir}/**/*no_context*generated*.jsonl", recursive=True) + ) + else: + files_model = sorted( + glob(f"{eval_model_dir}/**/*[!no_context]*generated*.jsonl", recursive=True) + ) + files_model = {file.split("/")[-1]: file for file in files_model} + + overlap_files = set(files_checkpoint) & set(files_model) + if not overlap_files: + raise ValueError( + "No overlapping files found between the evaluation checkpoint and model directory.\n" + f"Checkpoint files: {files_checkpoint.keys()}\n" + f"Model files: {files_model.keys()}" + ) + + print( + f"Generating results from {len(overlap_files)} overlapping files.\n" + + f"{overlap_files}" + ) + + out = { + "metadata": { + "source_path": checkpoint_name, + "custom_fields_schema": [ + {"name": "question", "type": "text"}, + {"name": "label", "type": "text"}, + {"name": "is_base_model_correct", "type": "category"}, + {"name": "is_checkpoint_model_correct", "type": "category"}, + {"name": "ctx_ids_len", "type": "number"}, + {"name": "query_ids_len", "type": "number"}, + ], + }, + "models": [{"name": checkpoint_name}, {"name": model_dir}], + "examples": [], + } + for file in overlap_files: + df_checkpoint = pd.read_json(path_or_buf=files_checkpoint[file], lines=True) + print(df_checkpoint.columns) + score_name = [col for col in df_checkpoint.columns if col.endswith("_score")] + assert len(score_name) == 1, ( + "Expected exactly one score column. Got: {score_name}" + ) + + df_model = pd.read_json(path_or_buf=files_model[file], lines=True) + print(df_model) + score_name = [col for col in df_model.columns if col.endswith("_score")] + assert len(score_name) == 1, ( + "Expected exactly one score column. Got: {score_name}" + ) + + score_name = score_name[0] + + # df_checkpoint.set_index("input", inplace=True) + df_checkpoint.drop(columns=["label"], inplace=True) + # df_model.set_index("input", inplace=True) + df = df_checkpoint.join(df_model, lsuffix="_checkpoint", rsuffix="_model") + # print(df.head()) + # print(df.tail()) + # print(df.columns) + # sample = df.iloc[0] + # print(sample["context"]) + # print(sample["input_checkpoint"]) + # print() + # print(sample["input_model"]) + df.drop(columns=["input_model"], inplace=True) + # print(f"{sample['generated_checkpoint']=}") + # print(f"{sample['generated_model']=}") + + for sample in df.itertuples(): + score_a = getattr(sample, f"{score_name}_checkpoint") + score_b = getattr(sample, f"{score_name}_model") + out["examples"].append( + { + "input_text": sample.context, + "tags": [file], + "output_text_a": sample.generated_checkpoint, + "output_text_b": sample.generated_model, + "score": score_a - score_b, + "individual_rater_scores": [], + "custom_fields": { + "question": sample.input_checkpoint, + "label": sample.label, + "is_base_model_correct": score_b >= 0.5, + "is_checkpoint_model_correct": score_a >= 0.5, + "ctx_ids_len": sample.ctx_ids_len, + "query_ids_len": sample.input_ids_len_checkpoint, + }, + } + ) + + # break + # # print(json_checkpoint[0]) + # # print() + # # print(json_model[0]) + return out + + +def build_confusion_matrix( + comparator_json: dict, + model_a_name: str = "checkpoint", + model_b_name: str = "model", +) -> pd.DataFrame: + """ + Build a confusion matrix from the comparator JSON. + """ + examples = comparator_json["examples"] + df = pd.DataFrame(examples) + df["is_correct_a"] = df["custom_fields"].apply( + lambda x: x.get("is_checkpoint_model_correct") + ) + df["is_correct_b"] = df["custom_fields"].apply( + lambda x: x.get("is_base_model_correct") + ) + + confusion_matrix = pd.crosstab( + df["is_correct_a"], + df["is_correct_b"], + rownames=[model_a_name], + colnames=[model_b_name], + margins=True, + margins_name="Total", + ) + return confusion_matrix + + +if __name__ == "__main__": + # Example usage + import json + + eval_checkpoint_dir = "train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/eval-results-105000" + model_dir = "google/gemma-2-2b-it" + comparator_json = convert_to_llm_comparator_format(eval_checkpoint_dir, model_dir) + confusion_matrix = build_confusion_matrix(comparator_json) + print(confusion_matrix) + # Save the comparator JSON to a file + with open(f"{eval_checkpoint_dir}/comparator.json", "w") as f: + json.dump(comparator_json, f)