first commit

Codes/sharedlibs/helperFunction/dataProcessing.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+"""
+clean_important_heads_outputs.py
+================================
+This script recursively traverses a result directory and processes each
+`important_heads.json` file by trimming the "assistant" part from the
+"output" field inside each response, keeping only the actual model output.
+"""
+
+import os
+import json
+from tqdm import tqdm
+
+def extract_assistant_only(output_text):
+    """Keep only the part after 'assistant' if present."""
+    if "assistant" in output_text:
+        return output_text.split("assistant", 1)[-1].strip()
+    else:
+        return output_text.strip()
+
+def process_json_file(file_path):
+    """Process a single important_heads.json file."""
+    with open(file_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    if "responses" not in data:
+        print(f"⚠️ Skipping file {file_path}: missing 'responses' field.")
+        return
+
+    for resp in data["responses"]:
+        if "output" in resp:
+            resp["output"] = extract_assistant_only(resp["output"])
+
+    # Overwrite the original file
+    with open(file_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
+
+def traverse_and_process(root_dir):
+    """Recursively traverse the directory and process all important_heads.json files."""
+    for dirpath, dirnames, filenames in os.walk(root_dir):
+        for filename in filenames:
+            if filename == "important_heads.json":
+                file_path = os.path.join(dirpath, filename)
+                process_json_file(file_path)
+
+if __name__ == "__main__":
+    # Replace with your actual root directory, e.g., "results"
+    root_directory = "results"
+    traverse_and_process(root_directory)
+    print("✅ All files have been processed.")