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