177 lines
6.7 KiB
Python
177 lines
6.7 KiB
Python
# 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()
|