51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
# -*- 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.")
|