Upload dataset_patcher.py

dataset_patcher.py

@@ -0,0 +1,346 @@
+import json
+import os
+import sys
+from collections import Counter
+from typing import Any, Dict, List, Optional, Tuple
+
+try:
+    import pandas as pd
+    import numpy as np
+except ImportError:
+    pd = None
+    np = None
+
+USER_ROLES = {"human", "user", "prompter"}
+ASSISTANT_ROLES = {"gpt", "assistant"}
+SYSTEM_ROLES = {"system"}
+
+
+def to_python_list(val: Any) -> Any:
+    """Safely converts numpy arrays / pandas series / iterables to native python list without truth checks."""
+    if val is None:
+        return None
+    if isinstance(val, (list, tuple)):
+        return list(val)
+    if np is not None and isinstance(val, np.ndarray):
+        return val.tolist()
+    if hasattr(val, "tolist"):
+        return val.tolist()
+    return val
+
+
+def inspect_entry(
+    entry: Dict[str, Any],
+) -> Tuple[bool, List[str], Optional[List[Dict[str, str]]]]:
+    errors = []
+
+    # Safely retrieve conversations without using `or` operator across arrays
+    raw_conv = entry.get("conversations")
+    if raw_conv is None or (isinstance(raw_conv, float) and pd.isna(raw_conv)):
+        raw_conv = entry.get("messages")
+
+    conv = to_python_list(raw_conv)
+
+    # Handle stringified JSON inside parquet columns
+    if isinstance(conv, str):
+        try:
+            conv = json.loads(conv)
+            conv = to_python_list(conv)
+        except Exception:
+            return (
+                False,
+                ["Field 'conversations'/'messages' contains invalid JSON string."],
+                None,
+            )
+
+    if conv is None:
+        return False, ["Missing 'conversations' or 'messages' key."], None
+    if not isinstance(conv, list):
+        return False, ["Conversation is not a list/array."], None
+    if len(conv) == 0:
+        return False, ["Conversation list is empty."], None
+
+    normalized_turns = []
+
+    for idx, turn in enumerate(conv):
+        # Handle stringified sub-elements if any
+        if isinstance(turn, str):
+            try:
+                turn = json.loads(turn)
+            except Exception:
+                errors.append(f"Turn #{idx} is not a valid dictionary or JSON.")
+                continue
+
+        if not isinstance(turn, dict):
+            errors.append(f"Turn #{idx} is not an object/dictionary.")
+            continue
+
+        raw_role = str(turn.get("from") or turn.get("role") or "").strip()
+        raw_content = str(turn.get("value") or turn.get("content") or "").strip()
+
+        if not raw_role:
+            errors.append(f"Turn #{idx} is missing role identifier.")
+        if not raw_content:
+            errors.append(f"Turn #{idx} has empty text content.")
+
+        role_lower = raw_role.lower()
+        canonical_role = None
+        if role_lower in USER_ROLES:
+            canonical_role = "human"
+        elif role_lower in ASSISTANT_ROLES:
+            canonical_role = "gpt"
+        elif role_lower in SYSTEM_ROLES:
+            canonical_role = "system"
+        else:
+            errors.append(f"Turn #{idx} has unrecognized role: '{raw_role}'.")
+
+        role_key = "from" if "from" in turn else "role"
+        val_key = "value" if "value" in turn else "content"
+
+        if raw_content and canonical_role:
+            normalized_turns.append(
+                {role_key: canonical_role, val_key: raw_content}
+            )
+
+    if not normalized_turns:
+        return False, errors or ["No valid text turns found."], None
+
+    last_turn_role = (
+        normalized_turns[-1].get("from") or normalized_turns[-1].get("role")
+    )
+    if last_turn_role != "gpt":
+        errors.append(f"Last turn is '{last_turn_role}' (must end with 'gpt').")
+
+    has_human = any(
+        (t.get("from") or t.get("role")) == "human" for t in normalized_turns
+    )
+    has_gpt = any(
+        (t.get("from") or t.get("role")) == "gpt" for t in normalized_turns
+    )
+
+    if not has_human:
+        errors.append("Missing at least one user/human turn.")
+    if not has_gpt:
+        errors.append("Missing at least one assistant/gpt turn.")
+
+    consecutive_dupes = False
+    for i in range(len(normalized_turns) - 1):
+        r1 = (
+            normalized_turns[i].get("from") or normalized_turns[i].get("role")
+        )
+        r2 = (
+            normalized_turns[i + 1].get("from")
+            or normalized_turns[i + 1].get("role")
+        )
+        if r1 == r2 and r1 != "system":
+            consecutive_dupes = True
+            break
+    if consecutive_dupes:
+        errors.append("Contains consecutive turns with the same role.")
+
+    repaired = None
+    if errors:
+        repaired = repair_conversation(normalized_turns)
+
+    is_valid = len(errors) == 0
+    return is_valid, errors, repaired
+
+
+def repair_conversation(
+    turns: List[Dict[str, str]],
+) -> Optional[List[Dict[str, str]]]:
+    if not turns:
+        return None
+
+    merged: List[Dict[str, str]] = []
+    for turn in turns:
+        role_k = "from" if "from" in turn else "role"
+        val_k = "value" if "value" in turn else "content"
+        curr_role = turn[role_k]
+        curr_val = turn[val_k]
+
+        if (
+            merged
+            and (merged[-1].get("from") or merged[-1].get("role")) == curr_role
+        ):
+            prev_val_k = "value" if "value" in merged[-1] else "content"
+            merged[-1][prev_val_k] += f"\n\n{curr_val}"
+        else:
+            merged.append({role_k: curr_role, val_k: curr_val})
+
+    while (
+        merged and (merged[-1].get("from") or merged[-1].get("role")) != "gpt"
+    ):
+        merged.pop()
+
+    has_human = any(
+        (t.get("from") or t.get("role")) == "human" for t in merged
+    )
+    has_gpt = any((t.get("from") or t.get("role")) == "gpt" for t in merged)
+
+    if merged and has_human and has_gpt:
+        return merged
+    return None
+
+
+def load_file(path: str) -> Tuple[List[Dict[str, Any]], str]:
+    ext = os.path.splitext(path)[1].lower()
+
+    if ext == ".parquet":
+        if pd is None:
+            raise ImportError(
+                "Reading .parquet requires pandas and pyarrow. Run: pip install pandas pyarrow"
+            )
+        df = pd.read_parquet(path)
+        return df.to_dict(orient="records"), "parquet"
+
+    with open(path, "r", encoding="utf-8") as f:
+        first_char = f.read(1)
+        f.seek(0)
+        if first_char == "[":
+            return json.load(f), "json"
+        else:
+            records = []
+            for line in f:
+                line = line.strip()
+                if line:
+                    records.append(json.loads(line))
+            return records, "jsonl"
+
+
+def main():
+    if len(sys.argv) < 2:
+        input_path = input(
+            "Enter path to JSON, JSONL, or Parquet dataset: "
+        ).strip()
+    else:
+        input_path = sys.argv[1]
+
+    input_path = input_path.strip("\"'")
+
+    if not os.path.isfile(input_path):
+        print(f"File not found: {input_path}")
+        sys.exit(1)
+
+    print(f"\n📂 Loading {input_path} into memory...")
+    try:
+        records, _ = load_file(input_path)
+    except Exception as e:
+        print(f"Error reading file: {e}")
+        sys.exit(1)
+
+    total_count = len(records)
+    print(f"✅ Loaded {total_count} records. Starting audit...\n")
+
+    valid_indices = []
+    invalid_data = []
+    error_counter = Counter()
+
+    for idx, record in enumerate(records):
+        is_valid, errors, repaired = inspect_entry(record)
+        if is_valid:
+            valid_indices.append(idx)
+        else:
+            invalid_data.append((idx, record, errors, repaired))
+            for err in errors:
+                error_counter[err] += 1
+
+    invalid_count = len(invalid_data)
+    valid_count = len(valid_indices)
+
+    print("=" * 60)
+    print("📊 AUDIT RESULTS SUMMARY")
+    print("=" * 60)
+    print(f"Total records audited : {total_count}")
+    print(f"Clean records         : {valid_count} ({(valid_count/total_count)*100:.2f}%)")
+    print(f"Malformed records     : {invalid_count} ({(invalid_count/total_count)*100:.2f}%)")
+    print("\nError Breakdown:")
+    for err, cnt in error_counter.most_common():
+        print(f"  • [{cnt} occurrences] {err}")
+    print("=" * 60)
+
+    if invalid_count == 0:
+        print("\n✨ No issues detected. Your dataset is 100% compliant.")
+        return
+
+    preview_limit = min(3, invalid_count)
+    print(f"\n🔍 PREVIEWING FIRST {preview_limit} MALFORMED ENTRIES:")
+    for i in range(preview_limit):
+        orig_idx, rec, errs, rep = invalid_data[i]
+        print(f"\n--- [Record #{orig_idx}] ---")
+        print(f"Issues Detected: {errs}")
+        raw_conv = rec.get("conversations")
+        if raw_conv is None:
+            raw_conv = rec.get("messages")
+        conv_preview = str(to_python_list(raw_conv))[:250]
+        print(f"Content: {conv_preview}...")
+        print(f"Repairable: {'Yes' if rep is not None else 'No'}")
+
+    print("\n" + "=" * 60)
+    print("🛠️ RESOLUTION OPTIONS:")
+    print("  [1] DELETE malformed records (keep only the 100% clean ones).")
+    print("  [2] REPAIR what is recoverable (drop only unfixable entries).")
+    print("  [3] CANCEL and make no changes.")
+    print("=" * 60)
+
+    choice = ""
+    while choice not in ["1", "2", "3"]:
+        choice = input("Select an option [1/2/3]: ").strip()
+
+    if choice == "3":
+        print("\nAborted. No changes written.")
+        sys.exit(0)
+
+    output_records = []
+    if choice == "1":
+        output_records = [records[i] for i in valid_indices]
+    elif choice == "2":
+        output_records = [records[i] for i in valid_indices]
+        repaired_success = 0
+        unrepairable_dropped = 0
+        for orig_idx, rec, errs, rep in invalid_data:
+            if rep is not None:
+                target_key = (
+                    "conversations" if "conversations" in rec else "messages"
+                )
+                rec[target_key] = rep
+                output_records.append(rec)
+                repaired_success += 1
+            else:
+                unrepairable_dropped += 1
+        print(f"\n  • Successfully repaired : {repaired_success}")
+        print(f"  • Unrepairable & dropped: {unrepairable_dropped}")
+
+    # Standardize output: ensure all array elements inside each record are native Python lists
+    for rec in output_records:
+        for k, v in list(rec.items()):
+            rec[k] = to_python_list(v)
+
+    base = os.path.splitext(input_path)[0]
+    default_out = f"{base}_clean.jsonl"
+    out_path = input(
+        f"\nEnter output path [Default: {default_out}]: "
+    ).strip().strip("\"'")
+    if not out_path:
+        out_path = default_out
+
+    print(f"\n💾 Writing {len(output_records)} records to {out_path}...")
+    out_ext = os.path.splitext(out_path)[1].lower()
+
+    if out_ext == ".parquet":
+        df_out = pd.DataFrame(output_records)
+        df_out.to_parquet(out_path, index=False)
+    elif out_ext == ".json":
+        with open(out_path, "w", encoding="utf-8") as f:
+            json.dump(output_records, f, ensure_ascii=False, indent=2)
+    else:
+        # Default to jsonl
+        with open(out_path, "w", encoding="utf-8") as f:
+            for row in output_records:
+                f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+    print(f"✅ Finished! Saved to: {os.path.abspath(out_path)}\n")
+
+
+if __name__ == "__main__":
+    main()