first commit

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

@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict, List
+import random
+
+
+def user(content: str) -> List[Dict[str, str]]:
+    return [{"role": "user", "content": content}]
+
+
+def tool(content: str) -> List[Dict[str, str]]:
+    # Use "tool" role as requested
+    return [{"role": "tool", "content": content}]
+
+
+def assistant(content: str) -> List[Dict[str, str]]:
+    return [{"role": "assistant", "content": content}]
+
+
+def suffile(parts: List[str]) -> str:
+    # You didn't specify the exact formatting; this is a simple, deterministic join.
+    # Change separator if you need (e.g., "\n\n", special tokens, etc.)
+    return "\n".join(parts)
+
+
+def read_json(path: Path) -> Any:
+    with path.open("r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def write_jsonl(path: Path, rows: List[Dict[str, Any]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        for row in rows:
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+
+def build_candidates(datasets: List[Dict[str, Any]], rng: random.Random) -> List[List[Dict[str, str]]]:
+    """
+    dataset_message_candidate: a list of message sequences (each sequence is a list of message dicts)
+    Built per your 7-variant recipe.
+    """
+    n = len(datasets)
+    if n < 2:
+        raise ValueError("Need at least 2 records to sample 'other data randomly' as requested.")
+
+    dataset_message_candidate: List[List[Dict[str, str]]] = []
+
+    for i, data in enumerate(datasets):
+        # from current data, we always take injection (as in your text)
+        inj_cur = data.get("injection", "")
+        # (you also said extract inst/cont/ans from each data, but then you override inst/cont/ans by sampling;
+        # we keep the extraction below for completeness, but we follow your sampling behavior.)
+        _inst_cur = data.get("instruction", "")
+        _cont_cur = data.get("input", "")
+        _ans_cur = data.get("output", "")
+
+        # Create 7 independent lists (NOT [[]]*7 which aliases)
+        message_list: List[List[Dict[str, str]]] = [[] for _ in range(7)]
+
+        # --- "2nd run": sample inst/cont/ans from OTHER data (paired)
+        j = rng.randrange(n - 1)
+        if j >= i:
+            j += 1
+        other1 = datasets[j]
+        inst_o1 = other1.get("instruction", "")
+        cont_o1 = other1.get("input", "")
+        ans_o1 = other1.get("output", "")
+
+        inst_wrapped = f"<inst>{inst_o1}</inst>"
+        inj_wrapped = f"<inst>{inj_cur}</inst>"
+
+        # message_list[0] += user(inst) + tool(cont) + assistant(ans)
+        message_list[0].extend(user(inst_wrapped))
+        message_list[0].extend(tool(cont_o1))
+        message_list[0].extend(assistant(ans_o1))
+
+        # message_list[1] += user(inst + cont) + assistant(ans)
+        message_list[1].extend(user(inst_wrapped + cont_o1))
+        message_list[1].extend(assistant(ans_o1))
+
+        # message_list[2] += user(inst + cont + inj) + assistant(ans)
+        message_list[2].extend(user(inst_wrapped + cont_o1 + inj_wrapped))
+        message_list[2].extend(assistant(ans_o1))
+
+        # message_list[3] += user(inst) + tool(cont + inj) + assistant(ans)
+        message_list[3].extend(user(inst_wrapped))
+        message_list[3].extend(tool(cont_o1 + inj_wrapped))
+        message_list[3].extend(assistant(ans_o1))
+
+        # --- Now sample inst2/cont2 from OTHER data (paired)
+        k = rng.randrange(n - 1)
+        if k >= i:
+            k += 1
+        other2 = datasets[k]
+        inst2 = other2.get("instruction", "")
+        cont2 = other2.get("input", "")
+
+        # NOTE: your pseudocode says:
+        # inst , inj= f"<inst>{inst2}</data>", f"<inst>{cont2}</data>"
+        # i.e. "inj" becomes wrapped cont2 (yes, weird, but we follow it).
+        inst2_wrapped = f"<inst>{inst2}</inst>"
+        cont2_wrapped = cont2
+
+        # You reference cont from earlier; in your pseudocode it's `cont` from the first sampled other1.
+        cont_base = cont_o1
+
+        # message_list[4] += user(inst) + tool(suffile([cont, inst2])) + assistant(ans)
+        message_list[4].extend(user(inst2_wrapped))
+        message_list[4].extend(tool(suffile([cont_base, inst2])))
+        message_list[4].extend(assistant(ans_o1))
+
+        # message_list[5] += user(inst) + tool(suffile([cont, cont2])) + assistant(ans)
+        message_list[5].extend(user(inst2_wrapped))
+        message_list[5].extend(tool(suffile([cont_base, cont2])))
+        message_list[5].extend(assistant(ans_o1))
+
+        # message_list[6] += user(inst) + tool(suffile([cont, inst2, const2])) + assistant(ans)
+        # Your pseudocode has a typo "const2"; interpret as cont2.
+        message_list[6].extend(user(inst2_wrapped))
+        message_list[6].extend(tool(suffile([cont_base, inst2, cont2])))
+        message_list[6].extend(assistant(ans_o1))
+
+        dataset_message_candidate.extend(message_list)
+
+    return dataset_message_candidate
+
+
+def build_output_dataset(
+    datasets: List[Dict[str, Any]],
+    dataset_message_candidate: List[List[Dict[str, str]]],
+    rng: random.Random,
+) -> List[Dict[str, Any]]:
+    """
+    output_dataset: per your pseudocode, for each record in datasets we build one temp_message_list by
+    concatenating random(1,10) candidates; after each concat, remove last assistant(ans).
+    We store each produced conversation as a JSONL row: {"messages": [...]}
+    """
+    output_rows: List[Dict[str, Any]] = []
+
+    for candidate in dataset_message_candidate:
+        temp_message_list: List[Dict[str, str]] = []
+
+        # random(1,10) -> interpret as randint(1, 9) because Python range(1,10) yields 1..9
+        for _i in range(rng.randrange(1, 2)):
+            if _i == 0:
+                cand = candidate
+            else:
+                cand = rng.choice(dataset_message_candidate)
+            temp_message_list.extend(cand)
+
+        # remove last assistant(ans)
+        if temp_message_list and temp_message_list[-1].get("role") == "assistant":
+            temp_message_list.pop()
+
+        output_rows.append({"messages": temp_message_list})
+
+    return output_rows
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--orig-inj-datapath", required=True, help="Path to original JSON dataset (list of dicts).")
+    ap.add_argument("--output-path", required=True, help="Output JSONL path.")
+    args = ap.parse_args()
+
+    in_path = Path(args.orig_inj_datapath)
+    out_path = Path(args.output_path)
+
+    datasets = read_json(in_path)
+    if not isinstance(datasets, list):
+        raise ValueError("Input JSON must be a list of records (dict).")
+
+    # Extract fields as you requested (even though later sampling uses 'other data')
+    # This also sanity-checks schema early.
+    for idx, d in enumerate(datasets):
+        if not isinstance(d, dict):
+            raise ValueError(f"Record {idx} is not a dict.")
+        for key in ("instruction", "input", "output", "injection"):
+            if key not in d:
+                raise ValueError(f"Record {idx} missing required key: {key}")
+
+    rng = random.Random(42)
+
+    dataset_message_candidate = build_candidates(datasets, rng)
+    print(len(dataset_message_candidate))
+    output_rows = build_output_dataset(datasets, dataset_message_candidate, rng)
+
+    write_jsonl(out_path, output_rows)
+    print(f"Wrote {len(output_rows)} conversations to {out_path}")
+
+
+if __name__ == "__main__":
+    main()

Codes/3-1_model_training_preprocess/inj-likechen/generate_training_dataset.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env sh
+set -eu
+
+python3 ../3-1_model_training_preprocess/inj-likechen/generate_training_dataset.py \
+  --orig-inj-datapath /data/local/hujk/IPIBench/topicattack/data/crafted_instruction_data_tri_injection_qa.json \
+  --output-path ../3-1_model_training_preprocess/inj-likechen/crafted_instruction_data_tri_injection_qa.jsonl

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()

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

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+# short_training_dataset_prompt.sh
+set -euo pipefail
+
+python3 generate_training_dataset2.py \
+  --orig-inj-datapath "../1_raw_dataset/topicattack/data/crafted_instruction_data_squad_injection_qa_train.json" \
+  --orig-inj-datapath-topicattack "../1_raw_dataset/topicattack/data/crafted_instruction_data_squad_conversation_attack_complete_train.json" \
+  --output-path "crafted_instruction_data_squad_conversation_attack_complete_train.jsonl"

Codes/3-1_model_training_preprocess/inj-likechen/generate_training_dataset_prompt.txt

@@ -0,0 +1,27 @@
+write a generate_training_dataset.py, and correspond shell script to foll parameters. --orig-inj-datapath (use /data/local/hujk/IPIBench/topicattack/data/ crafted_instruction_data_tri_injection_qa.json in sh), for each data, it extract inst=data["instruction"] , cont=data["input"], ans=data["output"], inj=data["injection"], and generate a new dataset and save to --output-path (use crafted_instruction_data_tri_injection_qa.jsonl in sh), and it create a list of message in this format. user() means {"role":"user", "content": content}, tool/assistent so on. create rng with seed 42, and use it in following code
+
+create rng with seed 42, and use it in following code
+
+dataset_message_candidate = []
+for each data in datasets:
+    message_list = [[]]*7
+    inst, cont, ans = select from other data randomly from 2nd run, inst/cont are paired, a.k.a in same data
+    inst , inj= f"<data>{inst}</data>", f"<data>{inj}</data>" #warp <data></data> to inst and inj
+    message_list[0] += user(inst) + tool(cont) + assistant(ans)
+    message_list[1] += user(inst + cont) + assistant(ans)
+    message_list[2] += user(inst + cont + inj) + assistant(ans)
+    message_list[3] += user(inst) + tool(cont + inj) + assistant(ans)
+    inst2, cont2 = select from other data randomly, inst2/cont2 are paired, a.k.a in same data
+    inst , inj= f"<data>{inst2}</data>", f"<data>{cont2}</data>"
+    message_list[4] += user(inst) + tool(suffile([cont, inst2])) + assistant(ans)
+    message_list[5] += user(inst) + tool(suffile([cont, cont2])) + assistant(ans)
+    message_list[6] += user(inst) + tool(suffile([cont, inst2, const2])) + assistant(ans)
+    dataset_message_candidate += message_list
+
+output_dataset = []
+for _ in range(len(datasets)):
+    temp_message_list = []
+    for i in range(random(1,10)):
+        temp_message_list += random choose(dataset_message_candidate)
+        remove last assistant(ans)
+    output_dataset += temp_message_list

Codes/3-1_model_training_preprocess/inj-likechen/short_training_dataset_prompt.txt

@@ -0,0 +1,34 @@
+write a short_training_dataset_prompt.py, and correspond short_training_dataset_prompt.sh to fill parameters.
+--orig-inj-datapath (use /data/local/hujk/IPIBench/topicattack/data/ crafted_instruction_data_tri_injection_qa.json in sh)
+for each data, it extract inst=data["instruction"] , cont=data["input"], ans=data["output"], inj=data["injection"], and generate a new dataset and save to 
+--output-path (use short_crafted_instruction_data_tri_injection_qa.jsonl in sh), and it create a list of message in this format:
+user() means {"role":"user", "content": content}, tool/assistent so on. the final dataset is a list for each line, [msg1, msg 2...]
+create rng with seed 42, and use it in following code
+
+dataset_message_candidate = []
+for each data in datasets:
+    message_list = [
+    inst, cont, ans = select from other data randomly from 2nd run, inst/cont are paired, a.k.a in same data
+    inst , inj= f"<inst>{inst}</inst>", f"<inst>{inj}</inst>" #warp <inst></inst> to inst and inj, don't touch content
+    # inst + inj combination
+    message_list += [[user(inst) + tool(cont) + assistant(ans)]]
+    message_list += [[user(inst) + tool(cont + inj) + assistant(ans)]]
+    message_list += [[user(inst) + tool(inj + cont) + assistant(ans)]]
+    message_list += [[user(inst + cont) + assistant(ans)]]
+    message_list += [[user(cont + inst) + assistant(ans)]]
+    message_list += [[user(inst + cont + inj) + assistant(ans)]]
+    
+    # inst + inst2 combination
+    inst2, cont2 = select from other data randomly, inst2/cont2 are paired, a.k.a in same data
+    inst2 = f"<inst>{inst2}</inst>"  #warp <inst></inst> to inst2, don't touch content
+    message_list += [[user(inst) + tool(suffile([cont, inst2])) + assistant(ans)]]
+    message_list += [[user(inst) + tool(suffile([cont, inst2, const2])) + assistant(ans)]]
+    dataset_message_candidate += message_list
+
+output_dataset = []
+for _ in range(len(datasets)):
+    temp_message_list = []
+    for i in range(random(1,2)):
+        temp_message_list += random choose(dataset_message_candidate)
+        remove last assistant(ans)
+    output_dataset += temp_message_list

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

@@ -0,0 +1,4 @@
+import sys
+from pathlib import Path
+
+print(str(Path(__file__).resolve().parents[2]))