first commit

Codes/3-1_model_training_preprocess/inj-likechen/generate_training_dataset2.py

@@ -0,0 +1,176 @@
+# short_training_dataset_prompt.py
+import argparse
+import json
+import os
+import random
+import sys
+import copy
+
+from pathlib import Path
+from typing import Callable, Dict, List
+from typing import Any, Dict, List
+sys.path.append(str(Path(__file__).resolve().parents[2]) + "/code")
+
+from lib.attack_defense_tools import escape_separation, ignore, naive, none, suffix_attack, completion_real, completion_realtmp, completion_realcmb, model_completion_real, conv_attack
+ATTACK_MAP: Dict[str, Callable] = {
+    "none": none,
+    "naive": naive,
+    "ignore": ignore,
+    "escape_separation": escape_separation,
+    "suffix_attack": suffix_attack,
+    "completion_real": completion_real,
+    "completion_realtmp": completion_realtmp,
+    "completion_realcmb": completion_realcmb,
+    "model_completion_real": model_completion_real,
+    "conv_attack": conv_attack
+}
+
+USED_ATTACK_LIST = ["ignore","escape_separation","completion_real","completion_realcmb","conv_attack"]
+
+def _merge_topicattack_data(data: List[dict], topic_data: List[dict]) -> List[dict]:
+    if len(data) != len(topic_data):
+        raise ValueError(
+            f"TopicAttack data length mismatch: base={len(data)} topic={len(topic_data)}"
+        )
+    merged = []
+    for idx, (base_item, topic_item) in enumerate(zip(data, topic_data)):
+        if "injection" not in topic_item:
+            raise KeyError(f"Missing injection in topicattack item {idx}")
+        merged_item = copy.deepcopy(base_item)
+        merged_item["injection_topicattack"] = topic_item["injection"]
+        merged.append(merged_item)
+    return merged
+
+
+def _apply_attack(d_item: dict, attack: str, side: str) -> dict:
+    attack_fn = ATTACK_MAP.get(attack)
+    if attack_fn is None:
+        raise ValueError(f"Unsupported attack: {attack}")
+    if attack == "conv_attack":
+        d_item["injection"] = d_item["injection_topicattack"]
+    return attack_fn(d_item, side=side, model=None)
+
+def user(content: str) -> Dict[str, str]:
+    return {"role": "user", "content": content}
+
+
+def tool(content: str) -> Dict[str, str]:
+    # You requested "tool/assistant so on" and your pseudocode uses tool(cont).
+    # If your training stack expects "system" or "assistant" here, change role accordingly.
+    return {"role": "tool", "content": content}
+
+
+def assistant(content: str) -> Dict[str, str]:
+    return {"role": "assistant", "content": content}
+
+
+def shuffle_join(parts: List[str], rng: random.Random) -> str:
+    parts2 = list(parts)
+    rng.shuffle(parts2)
+    return "".join(parts2)
+
+
+def load_json(path: str) -> List[Dict[str, Any]]:
+    with open(path, "r", encoding="utf-8") as f:
+        obj = json.load(f)
+    if not isinstance(obj, list):
+        raise ValueError(f"Expected a JSON list at {path}, got {type(obj)}")
+    return obj
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--orig-inj-datapath",
+        required=True,
+        help="Path to original injection dataset JSON (a list of dicts).",
+    )
+    ap.add_argument(
+        "--orig-inj-datapath-topicattack",
+        required=True,
+        help="Path to original injection dataset JSON (a list of dicts).",
+    )
+    ap.add_argument(
+        "--output-path",
+        required=True,
+        help="Output path (.jsonl). Each line is a JSON list of messages.",
+    )
+    args = ap.parse_args()
+
+    rng = random.Random(42)
+
+    datasets_orig = load_json(args.orig_inj_datapath)
+    datasets_topic = load_json(args.orig_inj_datapath_topicattack)
+    
+    datasets = _merge_topicattack_data(datasets_orig,datasets_topic)
+    if len(datasets) < 2:
+        raise ValueError("Need at least 2 data items to sample 'other data' examples.")
+    dataset_message_candidate: List[List[Dict[str, str]]] = []
+
+    n = len(datasets)
+    for i, data in enumerate(datasets):
+        inst = data["instruction"]
+        cont = data["input"]
+        ans = data["output"]
+        cont_w_attacks = []
+        for attack in USED_ATTACK_LIST:
+            d_item = {"input": cont, "injection": data["injection"] ,"injection_topicattack": data["injection_topicattack"] ,"add_mark":"inst"}
+            inj_attack = _apply_attack(d_item,attack,side="end")
+            cont_w_attacks += [inj_attack["input"]]
+
+        # pick another example for inst2/cont2 (paired from same data)
+        # j = rng.randrange(n - 1)
+        # if j >= i:
+        #     j += 1
+        # other = datasets[j]
+        # inst2 = other["instruction"]
+        # cont2 = other["input"]
+
+        # wrap inst & inj with <inst>...</inst>, DO NOT touch content strings
+        # inst_wrapped = f"<inst>{inst}</inst>"
+        # inj_wrapped = f"<inst>{inj}</inst>"
+        # inst2_wrapped = f"<inst>{inst2}</inst>"
+
+        message_list: List[List[Dict[str, str]]] = []
+
+        # inst + inj combinations (tool() holds "content" in your pseudocode)
+        for c in cont_w_attacks:
+            message_list += [[user(inst), assistant(""), tool(c), assistant(ans)]]
+        #message_list += [[user(inst_wrapped), tool(inj_wrapped + cont), assistant(ans)]]
+        #message_list += [[user(inst_wrapped + cont), assistant(ans)]]
+        #message_list += [[user(cont + inst_wrapped), assistant(ans)]]
+        #message_list += [[user(inst_wrapped + cont + inj_wrapped), assistant(ans)]]
+
+        # # inst + inst2 combinations
+        # message_list += [[user(inst_wrapped), tool(shuffle_join([cont, inst2_wrapped], rng)), assistant(ans)]]
+        # # Your pseudocode had: suffile([cont, inst2, const2]) (typo const2 -> cont2).
+        # message_list += [[user(inst_wrapped), tool(shuffle_join([cont, inst2_wrapped, cont2], rng)), assistant(ans)]]
+
+        dataset_message_candidate.extend(message_list)
+
+    # Build output dataset:
+    # For each original datum, pick k in {1,2} candidates, concatenating into one "conversation" per line.
+    # We remove the last assistant only for intermediate candidates, keeping a final assistant at the end.
+    output_dataset: List[List[Dict[str, str]]] = []
+    for cand_idx in range(len(dataset_message_candidate)):
+        k = rng.randint(1, 1)  # random(1,2) in your note -> interpreted as inclusive {1,2}
+        convo: List[Dict[str, str]] = []
+        for t in range(k):
+            if t == 0:
+                cand = dataset_message_candidate[cand_idx]
+            else:
+                cand = rng.choice(dataset_message_candidate)
+            convo.extend(cand)
+            # remove last assistant for all but the final appended candidate
+        if convo[-1]["role"] == "assistant":
+            convo.pop()
+        output_dataset.append(convo)
+
+    os.makedirs(os.path.dirname(args.output_path) or ".", exist_ok=True)
+    with open(args.output_path, "w", encoding="utf-8") as f:
+        for convo in output_dataset:
+            f.write(json.dumps(convo, ensure_ascii=False) + "\n")
+
+
+if __name__ == "__main__":
+    main()