From e692df98e9083aa75dbfec732d6cea0807876abb Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 20 May 2025 14:37:56 +0000 Subject: [PATCH] add negative_nq + visualize no_context base model --- README.md | 13 ++- generate_negative_natural_questions.py | 44 ++++++++ src/ctx_to_lora/data_definitions.py | 11 +- src/ctx_to_lora/data_utils.py | 30 +++--- src/ctx_to_lora/eval.py | 11 +- src/ctx_to_lora/training_utils.py | 2 +- webui/app.py | 144 ++++++++++++++++++++++--- webui/templates/visualize.html | 74 ++++++++++++- 8 files changed, 283 insertions(+), 46 deletions(-) create mode 100644 generate_negative_natural_questions.py diff --git a/README.md b/README.md index af81de4..49df497 100644 --- a/README.md +++ b/README.md @@ -51,12 +51,15 @@ run python intx_sft.py configs/...yaml ... --from_pretrained_checkpoint=train_ou LongBench ```bash # generative -run python src/ctx_to_lora/eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin +run python src/ctx_to_lora/eval.py --checkpoint_path train_outputs/runs/.../pytorch_model.bin --datasets negative_nq triviaqa_retrieved hotpot_qa squad longbench_e - +-split test -# benchmark -cd LongBench/LongBench +# run python src/ctx_to_lora/eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin -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 +# # benchmark +# cd LongBench/LongBench -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 +# 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 ``` \ No newline at end of file diff --git a/generate_negative_natural_questions.py b/generate_negative_natural_questions.py new file mode 100644 index 0000000..428f141 --- /dev/null +++ b/generate_negative_natural_questions.py @@ -0,0 +1,44 @@ +import os +from tqdm import tqdm +from datasets import load_dataset, Dataset + + +if __name__ == "__main__": + ds = load_dataset("google-research-datasets/natural_questions") + ds = ds.shuffle(seed=42) + for split in ds: + out = [] + for i, sample in enumerate(tqdm(ds[split])): + ctx = " ".join( + [ + token + for is_html, token in zip( + sample["document"]["tokens"]["is_html"], + sample["document"]["tokens"]["token"], + ) + if not is_html + ] + ).strip() + if not ctx: + continue + answers = [] + for answer in sample["annotations"]["short_answers"]: + answers += answer["text"] + if not answers: + continue + answer = answers[0].strip() + if not answer: + continue + prompt = sample["question"]["text"].capitalize().strip() + " ?" + out.append(dict(prompt=prompt, context=ctx, answer=answer.capitalize())) + # if i >= 10: + # break + shifted_ctxs = [d["context"] for d in out] + shifted_ctxs = shifted_ctxs[1:] + shifted_ctxs[:1] + for d, shifted_ctx in zip(out, shifted_ctxs): + d["context"] = shifted_ctx + new_ds = Dataset.from_list(out) + print(new_ds) + save_dir = "data/raw_datasets/negative_natural_questions" + os.makedirs(save_dir, exist_ok=True) + new_ds.to_json(f"{save_dir}/{split}.jsonl") diff --git a/src/ctx_to_lora/data_definitions.py b/src/ctx_to_lora/data_definitions.py index 40d67fc..0efc52c 100644 --- a/src/ctx_to_lora/data_definitions.py +++ b/src/ctx_to_lora/data_definitions.py @@ -299,8 +299,14 @@ DS_KWARGS = { split="train", ), ), - # TODO: for negative training and evaluating - "natural_questions": ..., + # for negative training and evaluating + "negative_nq": dict( + test=dict( + path="json", + data_files="data/raw_datasets/negative_natural_questions/validation.jsonl", + split="train", + ), + ), } # LongBench kwargs @@ -341,6 +347,7 @@ CLOSED_QA_INTX_TEMPLATES = [ EVAL_INTX_TEMPLATES = { + "negative_nq": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", "triviaqa_retrieved": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", "hotpot_qa": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", "squad": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", diff --git a/src/ctx_to_lora/data_utils.py b/src/ctx_to_lora/data_utils.py index 2806fa2..38deb21 100644 --- a/src/ctx_to_lora/data_utils.py +++ b/src/ctx_to_lora/data_utils.py @@ -192,25 +192,19 @@ def get_preprocessing_fn( "response": sample["answers"][0], } - elif "natural_questions" in ds_name: - pass - # ctx = " ".join( - # [ - # token - # for is_html, token in zip( - # ds[0]["document"]["tokens"]["is_html"], - # ds[0]["document"]["tokens"]["token"], - # ) - # if not is_html - # ] - # ) - elif ds_name == "triviaqa_retrieved": - # maybe needed for training (negative sample for training) - # def f(sample): - # ctx = sample["entity_page"]["wiki_context"] - # if not ctx: - # return None + elif ds_name == "negative_nq": + def f(sample): + q = sample["prompt"] + prompt = closed_qa_prompting(q) if not is_eval else q + return { + "context": sample["context"], + "prompt": prompt, + "response": sample["answer"], + } + + elif ds_name == "triviaqa_retrieved": + # only used for eval def f(sample): return { "context": sample["context"], diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py index 50bb7e0..1290cd3 100644 --- a/src/ctx_to_lora/eval.py +++ b/src/ctx_to_lora/eval.py @@ -95,6 +95,7 @@ CLOSED_QA_DATASETS = { "longbench/musique", "squad", "triviaqa_retrieved", + "negative_nq", } for ds_name in list(CLOSED_QA_DATASETS): @@ -452,7 +453,13 @@ def evaluate( # if generative: # model = model.to(torch.bfloat16) else: - model = get_model(model_name_or_path, train=False, requires_grad=False) + model_kwargs = dict(attn_implementation="flash_attention_2") + model = get_model( + model_name_or_path, + train=False, + requires_grad=False, + model_kwargs=model_kwargs, + ) # NOTE: there is still some randomness in the eval result # despite all the deterministic settings if is_liger_kernel_available(): @@ -681,7 +688,7 @@ if __name__ == "__main__": args.output_dir = f"{run_dir}/eval-results-{cur_it}" args.logging_dir = f"{run_dir}/eval-results-{cur_it}" args.run_name = run_dir.split("/")[-1] - args.remove_context = True + args.remove_context = False # modulated model doesn't see ctx by default else: args = Namespace( model_name_or_path=cli_args.model_name_or_path, diff --git a/src/ctx_to_lora/training_utils.py b/src/ctx_to_lora/training_utils.py index 3e4ca2a..b76858d 100644 --- a/src/ctx_to_lora/training_utils.py +++ b/src/ctx_to_lora/training_utils.py @@ -124,7 +124,7 @@ def train_model( compute_metrics=compute_metrics, ) if is_modulated_model: - logger.info(f"Training with modulated model. Using CustomTrainer.") + logger.info("Training with modulated model. Using CustomTrainer.") trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef del training_args.gen_lora_l1_reg_coef diff --git a/webui/app.py b/webui/app.py index 4ccb8c3..22052d6 100644 --- a/webui/app.py +++ b/webui/app.py @@ -116,51 +116,134 @@ def get_run_data(run_path): return grouped_data -def get_generated_text_data(run_path): +def get_generated_text_data(run_path: str) -> dict: """ - Loads generated text data from all *_generated_text.jsonl files. + Loads generated text data from all *_generated_text.jsonl files from the current run's output. + It specifically excludes any files ending with _no_context_generated_text.jsonl, + as those are handled by get_base_model_generated_text. Also recursively searches for generated text files in subfolders. Args: - run_path: Path to the directory containing the .jsonl files. + run_path: Path to the directory containing the .jsonl files for the current run. Returns: - A dictionary containing the generated text data for each split found. + A dictionary containing the generated text data for each primary split found. """ generated_data = {} - files = sorted(glob(f"{run_path}/*_generated_text.jsonl")) - print(files) - # Find all *_generated_text.jsonl files in the current directory + # Find all *_generated_text.jsonl files in the current directory, excluding _no_context_ variants + files = sorted( + [ + f + for f in glob(os.path.join(run_path, "*_generated_text.jsonl")) + if "_no_context_generated_text.jsonl" not in os.path.basename(f) + ] + ) + print(f"Processing generated text files from {run_path}: {files}") + for filename in files: - # Extract split name from filename (remove _generated_text.jsonl) - split = filename.split("/")[-1].split("_generated_text.jsonl")[0] + split = os.path.basename(filename).replace("_generated_text.jsonl", "") try: with open(filename) as f: - lines = f.readlines() + lines = f.readlines(100_000) data = [json.loads(line) for line in lines] generated_data[split] = data except FileNotFoundError: generated_data[split] = None + except json.JSONDecodeError as e: + print(f"Error decoding JSON from {filename}: {e}") + generated_data[split] = None # Recursively search for *_generated_text.jsonl files in subdirectories - for subdir in [d for d in glob(os.path.join(run_path, "*")) if os.path.isdir(d)]: - subdir_name = os.path.basename(subdir) - subdir_files = sorted(glob(f"{subdir}/*_generated_text.jsonl")) + for subdir_path in [ + d for d in glob(os.path.join(run_path, "*")) if os.path.isdir(d) + ]: + subdir_name = os.path.basename(subdir_path) + subdir_files = sorted( + [ + f + for f in glob(os.path.join(subdir_path, "*_generated_text.jsonl")) + if "_no_context_generated_text.jsonl" not in os.path.basename(f) + ] + ) + print( + f"Processing generated text files from subdirectory {subdir_path}: {subdir_files}" + ) for filename in subdir_files: - # Extract split name from filename and include subfolder name - base_split = filename.split("/")[-1].split("_generated_text.jsonl")[0] + base_split = os.path.basename(filename).replace("_generated_text.jsonl", "") + # Create a unique split key including the subdirectory name split = f"{subdir_name}_{base_split}" try: with open(filename) as f: - lines = f.readlines() + lines = f.readlines(100_000) data = [json.loads(line) for line in lines] generated_data[split] = data except FileNotFoundError: generated_data[split] = None - + except json.JSONDecodeError as e: + print(f"Error decoding JSON from {filename}: {e}") + generated_data[split] = None return generated_data +def get_base_model_generated_text(model_name, generated_data): + """ + Loads generated text data from the base model for comparison. + Also loads "no context" base model data if available. + + Args: + model_name: The name of the base model + generated_data: Dictionary of generated data from the fine-tuned model + + Returns: + Tuple: (base_model_data, base_model_no_context_data) + Dictionaries containing base model generated texts for matching splits + """ + if not model_name: + return {}, {} + + base_model_data = {} + base_model_no_context_data = {} + + # # Create normalized model name for directory lookup + # normalized_model_name = model_name.replace("/", "_") + + for split in generated_data: + # Skip if no data for this split + if not generated_data[split]: + base_model_data[split] = None + base_model_no_context_data[split] = None + continue + + # Construct path to the base model's output for this split + base_model_path = f"eval_results/{model_name}/{split}_generated_text.jsonl" + base_model_no_context_path = ( + f"eval_results/{model_name}/{split}_no_context_generated_text.jsonl" + ) + + try: + with open(base_model_path) as f: + lines = f.readlines(100_000) + data = [json.loads(line) for line in lines] + base_model_data[split] = data + except FileNotFoundError: + print(f"Base model output not found for {split} at {base_model_path}") + # Store None to indicate we tried but didn't find matching data + base_model_data[split] = None + + try: + with open(base_model_no_context_path) as f: + lines = f.readlines(100_000) + data = [json.loads(line) for line in lines] + base_model_no_context_data[split] = data + except FileNotFoundError: + print( + f"Base model (no context) output not found for {split} at {base_model_no_context_path}" + ) + base_model_no_context_data[split] = None + + return base_model_data, base_model_no_context_data + + def get_available_checkpoints(run): """ Finds all checkpoint directories in a run folder. @@ -292,6 +375,31 @@ def visualize(run): except FileNotFoundError: model_name = None + # If config.yaml doesn't have the model name, try args.yaml + if not model_name: + args_path = os.path.join(logdir, "args.yaml") + try: + with open(args_path) as f: + yaml.add_constructor("!", lambda loader, node: None) + yaml.add_multi_constructor( + "tag:yaml.org,2002:python/object", + lambda loader, suffix, node: None, + Loader=yaml.SafeLoader, + ) + args = yaml.safe_load(f) + model_name = args.get("model_name_or_path") + except FileNotFoundError: + model_name = None + + # Get base model outputs if we have the model name and modulated outputs + base_model_data = {} + base_model_no_context_data = {} + print(generated_data.keys()) + if model_name and generated_data: + base_model_data, base_model_no_context_data = get_base_model_generated_text( + model_name, generated_data + ) + return render_template( "visualize.html", run=run, @@ -299,6 +407,8 @@ def visualize(run): selected_eval_folder=selected_eval_folder, data=data, generated_data=generated_data, + base_model_data=base_model_data, + base_model_no_context_data=base_model_no_context_data, model_name=model_name, checkpoints=checkpoint_names, ) diff --git a/webui/templates/visualize.html b/webui/templates/visualize.html index 67fd98a..a81623e 100644 --- a/webui/templates/visualize.html +++ b/webui/templates/visualize.html @@ -452,6 +452,43 @@ background-color: #bbdefb; cursor: not-allowed; } + + .model-comparison { + display: flex; + flex-wrap: wrap; + gap: 20px; + margin-bottom: 15px; + } + + .model-output { + flex: 1; + min-width: 300px; + border: 1px solid #e0e0e0; + border-radius: 5px; + padding: 10px; + background-color: #f9f9f9; + } + + .model-output p { + margin: 0; + } + + .model-output span { + display: block; + padding: 10px; + background-color: white; + border: 1px solid #eee; + border-radius: 4px; + min-height: 80px; + max-height: 150px; + overflow-y: auto; + white-space: pre-wrap; + margin-top: 5px; + } + + .highlight-diff { + background-color: #fff8e1; + } {% endblock %} @@ -536,7 +573,22 @@ Context:

Input:

-

Generated:

+ + +
+
+

Modulated Model:

+
+
+

Base Model ({{ model_name }}): Not + available

+
+
+

Base Model (No Context - {{ model_name }}): Not available

+
+
+

Label:

@@ -632,10 +684,13 @@