Claude/serene einstein ej ask

README.md

@@ -155,6 +155,14 @@ Pick the script that matches your base model:
 | Meta-Llama-3-8B-Instruct | `bash ./scripts/llama8b/sep/drip_sep.sh` |
 | Mistral-7B-Instruct-v0.3 | `bash ./scripts/mistral7b/sep/drip_sep.sh` |
 
+Training merges the LoRA adapter into the base weights and saves a **full
+checkpoint** that evaluation can load directly. For checkpoints saved as adapters
+instead (QLoRA runs, or models trained earlier), merge them first:
+
+```bash
+python merge_lora.py --adapter_path <adapter_dir> --output_path <merged_dir>
+```
+
 ---
 
 ## Evaluation

data_generation/README.md

@@ -24,7 +24,7 @@ Run this script over the injected dataset.
 
 1. **Load** injected requests from `{name}_injected_diff_output.json`.
 2. **`process_chosen`** wraps input as `<instruction>...</instruction><start of data>...<end of data>`, strips the `INSISTENCE` marker, and queries the model with a system message enforcing data/instruction separation.
-   - If `--no_judge` is off, an LLM judge checks whether the probe leaked into the answer. On `"INCLUDED"`, it retries once with corrective feedback.
+   - If `--no_judge` is off, an LLM judge ([`tasktracker_judge_prompt2.txt`](./tasktracker_judge_prompt2.txt)) checks whether the model actually **used the accompanying text as data** (the probe text should remain present in the data, just not executed). It returns `INCLUDED` (data was used) or `NOT INCLUDED` (data was ignored/omitted); on `"NOT INCLUDED"` it retries once, asking the model to keep the probe as part of the data.
 3. **`process_rejected`** feeds the `injected_probe` as a developer message over `clean_input`, producing the non-compliant (injected) response.
 4. **`build_dpo_pairs`** joins chosen and rejected on `(instruction, input)`, then the result is shuffled and saved.
 

data_generation/check_sep_leakage.py

@@ -0,0 +1,113 @@
+"""Check that the SEP (or Alpaca) train and eval sets are disjoint at the
+content level — guards against train/test leakage that would inflate results.
+
+The DRIP training set (datasets/sep/train_dataset.json) and the SEP eval set
+(datasets/SEP_dataset.json) use different field names but the same underlying
+(system prompt, clean data) content. This script normalizes that content and
+reports any overlap. It is read-only and makes no changes.
+
+Usage:
+    python data_generation/check_sep_leakage.py \
+        --train datasets/sep/train_dataset.json \
+        --eval  datasets/SEP_dataset.json
+
+    # Alpaca injection split (override the field names):
+    python data_generation/check_sep_leakage.py \
+        --train datasets/alpaca_data.json \
+        --eval  datasets/davinci_003_outputs.json \
+        --sys-fields instruction --data-fields input
+
+Exit code is non-zero if any (system, clean-data) pair appears in both sets,
+so it can also be used as a CI assertion.
+"""
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+
+# Candidate field names (first match wins) for the two halves of the key.
+SYS_FIELDS = ["system_prompt_clean", "system_prompt", "instruction"]
+DATA_FIELDS = ["data_prompt_clean", "prompt_clean", "input", "data"]
+
+
+def _load(path: Path):
+    text = path.read_text()
+    stripped = text.lstrip()
+    if stripped.startswith("["):
+        return json.loads(text)
+    # JSONL fallback
+    return [json.loads(line) for line in text.splitlines() if line.strip()]
+
+
+def _norm(s) -> str:
+    return re.sub(r"\s+", " ", str(s or "").strip().lower())
+
+
+def _pick(item: dict, fields):
+    for f in fields:
+        if f in item and item[f] not in (None, ""):
+            return item[f]
+    return ""
+
+
+def _keys(items, sys_fields, data_fields):
+    sys_set, data_set, both_set = set(), set(), set()
+    for it in items:
+        s = _norm(_pick(it, sys_fields))
+        d = _norm(_pick(it, data_fields))
+        if s:
+            sys_set.add(s)
+        if d:
+            data_set.add(d)
+        if s or d:
+            both_set.add((s, d))
+    return sys_set, data_set, both_set
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--train", required=True, type=Path)
+    ap.add_argument("--eval", required=True, type=Path)
+    ap.add_argument("--sys-fields", nargs="+", default=SYS_FIELDS,
+                    help="Candidate field names for the system/instruction text.")
+    ap.add_argument("--data-fields", nargs="+", default=DATA_FIELDS,
+                    help="Candidate field names for the clean data/input text.")
+    ap.add_argument("--show", type=int, default=3, help="How many overlapping examples to print.")
+    args = ap.parse_args()
+
+    train, ev = _load(args.train), _load(args.eval)
+    print(f"train: {len(train)} items  |  eval: {len(ev)} items")
+
+    t_sys, t_data, t_both = _keys(train, args.sys_fields, args.data_fields)
+    e_sys, e_data, e_both = _keys(ev, args.sys_fields, args.data_fields)
+
+    both = t_both & e_both
+    data_only = t_data & e_data
+    sys_only = t_sys & e_sys
+
+    def pct(n, d):
+        return f"{100 * n / d:.1f}%" if d else "n/a"
+
+    print(f"\n(system, clean-data) pairs in BOTH : {len(both):5d}  "
+          f"({pct(len(both), len(e_both))} of eval)")
+    print(f"clean-data strings in BOTH          : {len(data_only):5d}  "
+          f"({pct(len(data_only), len(e_data))} of eval)")
+    print(f"system strings in BOTH              : {len(sys_only):5d}  "
+          f"({pct(len(sys_only), len(e_sys))} of eval)")
+
+    if both and args.show:
+        print(f"\n--- {min(args.show, len(both))} overlapping (system, data) examples ---")
+        for s, d in list(both)[:args.show]:
+            print(f"  SYS : {s[:100]}")
+            print(f"  DATA: {d[:100]}\n")
+
+    if both:
+        print("LEAKAGE DETECTED: eval items also appear in train (see above).")
+        sys.exit(1)
+    print("OK: no (system, clean-data) overlap between train and eval.")
+
+
+if __name__ == "__main__":
+    main()

data_generation/data_curation_drip.py

@@ -112,7 +112,13 @@ async def run_batch(processor: OptimizedAPIProcessor,
                     batch_size: int = 100,
                     starting_index: int = 0) -> List[Dict]:
     existing = jload(output_path) if os.path.exists(output_path) else []
-    remaining = requests[len(existing):]
+    # Resume by content key, not by count: dropped failures (None results) must
+    # not shift a length-based offset, which would skip unprocessed requests and
+    # reprocess completed ones. Key on (instruction, injected_input), matching the
+    # output records' (instruction, input).
+    done_keys = {(r["instruction"], r["input"]) for r in existing}
+    remaining = [req for req in requests
+                 if (req["instruction"], req["injected_input"]) not in done_keys]
 
     if not remaining:
         print("All requests already processed!")

data_generation/dpo_data_loader.py

@@ -27,7 +27,7 @@ def generate_training_data_dpo(data_dicts, prompt_dict_name, tokenizer):
 
 class PreferenceWithExpertDatasetHF:
     def __init__(self, data_path_list: List[str], tokenizer: PreTrainedTokenizer, attack: str):
-        prompt_dict_name, attacks = attack.split('_')
+        prompt_dict_name, attacks = attack.split('_', 1)
         prompts: List[str] = []
         chosens: List[str] = []
         rejecteds: List[str] = []
@@ -78,7 +78,13 @@ def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
 
         # 2) resp: chosen / rejected
         def enc(texts):
-            out = self.tokenizer(texts, padding=False, truncation=True, max_length=self.max_target_length, return_tensors=None)
+            # add_special_tokens=False: the prompt (tokenized separately, with BOS)
+            # is prepended to this response, so the response must NOT get its own
+            # BOS — otherwise a spurious BOS lands mid-sequence at the prompt/
+            # response boundary (train/inference mismatch). The explicit eos in the
+            # text is still tokenized to the eos id.
+            out = self.tokenizer(texts, padding=False, truncation=True, max_length=self.max_target_length,
+                                 add_special_tokens=False, return_tensors=None)
             ids = [torch.tensor(x, dtype=torch.long) for x in out["input_ids"]]
             return ids
 

data_generation/sft_data_loader.py

@@ -375,7 +375,7 @@ class SupervisedDataset(Dataset):
     def __init__(self, data_path_list: List[str], tokenizer, attack, frontend_delimiters, downsample=True):
         super(SupervisedDataset, self).__init__()
 
-        prompt_dict_name, attacks = attack.split('_')
+        prompt_dict_name, attacks = attack.split('_', 1)
 
         self.input_ids = []
         self.labels = []

merge_lora.py

@@ -0,0 +1,70 @@
+"""Merge a trained LoRA adapter into its base model and save a full checkpoint.
+
+Use this for checkpoints saved as adapters — e.g. QLoRA runs, or models trained
+before train_unified.py started merging at save time. It loads the model exactly
+the way evaluation does (testing/test.py `load_full_model`, load_as_adapter=True),
+so the merged result is guaranteed consistent with eval, then merges and saves a
+standalone checkpoint that the default eval path can load directly.
+
+Usage:
+    python merge_lora.py \
+        --adapter_path  out/Meta-Llama-3-8B-Instruct-TextTextText-sep-drip \
+        --output_path   out/Meta-Llama-3-8B-Instruct-TextTextText-sep-drip-merged
+
+    # base path / model class are inferred from the adapter path; override if needed:
+    python merge_lora.py --adapter_path <dir> --output_path <dir> \
+        --base_model_path meta-llama/Meta-Llama-3-8B-Instruct \
+        --customized_model_class LlamaForCausalLMDRIP
+"""
+
+import argparse
+
+from testing.test import load_full_model
+
+
+def main():
+    ap = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    ap.add_argument("--adapter_path", required=True,
+                    help="Directory of the trained LoRA adapter checkpoint.")
+    ap.add_argument("--output_path", required=True,
+                    help="Where to write the merged full checkpoint.")
+    ap.add_argument("--base_model_path", default=None,
+                    help="Base model (inferred from --adapter_path if omitted).")
+    ap.add_argument("--customized_model_class", default="",
+                    help="REGISTRY class key (auto-detected from the path if omitted).")
+    ap.add_argument("--device", default="auto",
+                    help="device_map for loading (e.g. 'auto', '0', 'cpu').")
+    args = ap.parse_args()
+
+    device_map = int(args.device) if args.device.isdigit() else args.device
+
+    # Same loading path as evaluation; load_as_adapter attaches the adapter onto
+    # the base model. The model class / delimiters are resolved internally.
+    model, tokenizer = load_full_model(
+        args.adapter_path,
+        customized_model_class=args.customized_model_class,
+        load_as_adapter=True,
+        base_model_path=args.base_model_path,
+        device_map=device_map,
+    )
+
+    if not hasattr(model, "merge_and_unload"):
+        raise SystemExit(
+            "Loaded model is not a PEFT adapter (nothing to merge). "
+            "Check --adapter_path / --customized_model_class."
+        )
+
+    print("Merging adapter into base weights ...")
+    merged = model.merge_and_unload()
+
+    merged.save_pretrained(args.output_path, safe_serialization=True)
+    tokenizer.save_pretrained(args.output_path)
+    # Persist the (DRIP) config — including delimiter ids — alongside the weights.
+    getattr(merged, "base_model", merged).config.save_pretrained(args.output_path)
+    print(f"Merged checkpoint saved to {args.output_path}")
+
+
+if __name__ == "__main__":
+    main()

testing/test.py

@@ -369,7 +369,9 @@ def form_llm_input(data, injection_method, fmt, apply_filter, defense, sample_id
     llm_input = injection_method(fmt) if injection_method is hackaprompt else []
     for i, d in enumerate(data):
         if injection_method is none:
-            llm_input.append(apply_testtime_defense(d, fmt, defense))
+            # deepcopy: apply_testtime_defense mutates the dict in place, so never
+            # pass the shared source `d` (would contaminate later attacks/runs).
+            llm_input.append(apply_testtime_defense(deepcopy(d), fmt, defense))
             continue
         if not d["input"] or injection_method is hackaprompt:
             continue

train_unified.py

@@ -665,14 +665,27 @@ def main(argv: Optional[List[str]] = None):
     trainer.train(resume_from_checkpoint=resolve_resume(training_args, trainer))
 
     # ---- Save ----
-    if trainer.is_world_process_zero():
-        logger.info("Training done. Saving...")
-    trainer.save_model(output_dir=training_args.output_dir)
+    # For LoRA runs, merge the adapter into the base weights so evaluation can
+    # load a full checkpoint directly (the default path in testing/test.py).
+    # QLoRA/4-bit adapters cannot be merged, so those are saved as adapters
+    # (load them with --load_as_adapter). AIR/full-finetune (non-PEFT) just
+    # saves the full model.
     trainer.save_state()
     if trainer.is_world_process_zero():
+        logger.info("Training done. Saving...")
+        to_save = trainer.accelerator.unwrap_model(trainer.model)
+        if bnb_config is None and hasattr(to_save, "merge_and_unload"):
+            try:
+                to_save = to_save.merge_and_unload()
+                logger.info("Merged LoRA adapter into base weights.")
+            except Exception as e:  # noqa: BLE001
+                logger.warning(f"merge_and_unload failed ({e}); saving adapter instead.")
+        to_save.save_pretrained(training_args.output_dir, safe_serialization=True)
         tokenizer.save_pretrained(training_args.output_dir)
-        base_config = getattr(trainer.model, "base_model", trainer.model).config
-        base_config.save_pretrained(training_args.output_dir)
+        # Persist the (DRIP) config — incl. delimiter ids — alongside the weights.
+        getattr(to_save, "base_model", to_save).config.save_pretrained(training_args.output_dir)
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
 
 
 if __name__ == "__main__":

training/trainer.py

@@ -294,12 +294,12 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
             chosen_logps, chosen_mean_logits, chosen_aux = self._seq_logps_moe(
                 self.model, inputs["chosen_input_ids"], inputs["chosen_attention_mask"],
                 inputs["chosen_expert_labels"], inputs["prompt_lens"],
-                return_logits=True, output_router_logits=False)
+                return_logits=True, output_router_logits=True)
 
             rejected_logps, rejected_mean_logits, rejected_aux = self._seq_logps_moe(
                 self.model, inputs["rejected_input_ids"], inputs["rejected_attention_mask"],
                 inputs["rejected_expert_labels"], inputs["prompt_lens"],
-                return_logits=True, output_router_logits=False)
+                return_logits=True, output_router_logits=True)
 
             # Ref forward: NO aux_loss needed (ref is frozen)
             with torch.no_grad():
@@ -365,7 +365,9 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
 
 
 class DPOTrainerAIR(DPOTrainerOurs):
-    def create_optimizer(self, special_params_list = ["intermediate_shifts"]):
+    def create_optimizer(self, special_params_list=None):
+        if special_params_list is None:
+            special_params_list = ["intermediate_shifts"]
         opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
         if self.optimizer is None:
             decay_parameters = self.get_decay_parameter_names(opt_model)