wehub-resource-sync
365 lines
16 KiB
Python
365 lines
16 KiB
Python
import json
|
|
import os
|
|
from tqdm import tqdm
|
|
from multiprocessing import Pool, cpu_count
|
|
import yaml
|
|
|
|
|
|
class DataProcessor:
|
|
def __init__(self, file_path, image_root, video_root):
|
|
self.file_path = file_path
|
|
self.image_root = image_root
|
|
self.data = None
|
|
self.video_root = video_root
|
|
self.load_data()
|
|
|
|
def load_data(self):
|
|
if self.file_path.endswith(".json"):
|
|
with open(self.file_path, "r") as f:
|
|
self.data = json.load(f)
|
|
elif self.file_path.endswith(".yaml"):
|
|
with open(self.file_path, "r") as f:
|
|
self.data = yaml.safe_load(f)
|
|
elif self.file_path.endswith(".jsonl"):
|
|
with open(self.file_path, "r") as f:
|
|
self.data = [json.loads(line) for line in f.readlines()]
|
|
else:
|
|
raise ValueError("Unsupported file format")
|
|
|
|
def load_json_data(self, json_path):
|
|
if json_path.endswith(".jsonl"):
|
|
cur_data_dict = []
|
|
with open(json_path, "r") as json_file:
|
|
for line in json_file:
|
|
cur_data_dict.append(json.loads(line.strip()))
|
|
return cur_data_dict
|
|
elif json_path.endswith(".json"):
|
|
with open(json_path, "r") as f:
|
|
return json.load(f)
|
|
else:
|
|
raise ValueError("Unsupported file format")
|
|
|
|
def check_image_existence(self, data):
|
|
if "image" in data:
|
|
if type(data["image"]) == list:
|
|
images = data["image"]
|
|
else:
|
|
images = [data["image"]]
|
|
|
|
for image in images:
|
|
full_image_path = os.path.join(self.image_root, image)
|
|
if not os.path.exists(full_image_path):
|
|
print(f"WARNING!!! {full_image_path} not exists !!!")
|
|
|
|
if "video" in data:
|
|
full_video_path = os.path.join(self.video_root, data["video"])
|
|
if not os.path.exists(full_video_path):
|
|
print(f"WARNING!!! {full_video_path} not exists !!!")
|
|
|
|
# if data["conversations"][0]["value"].count("<image>") > 1:
|
|
# print(f"WARNING!!! {data['conversations'][0]['value']} has more than one <image> !!!")
|
|
|
|
def check_item_structure(self, item):
|
|
if not all(key in item for key in ["conversations"]):
|
|
print(f"WARNING!!! Item {item.get('id', 'unknown')} is missing required fields!")
|
|
return False
|
|
|
|
conversations = item["conversations"]
|
|
if not isinstance(conversations, list) or len(conversations) < 2 or len(conversations) % 2 != 0:
|
|
print(f"WARNING!!! Item {item['id']} has invalid conversations structure!")
|
|
return False
|
|
|
|
for i, conv in enumerate(conversations):
|
|
if not all(key in conv for key in ["from", "value"]):
|
|
print(f"WARNING!!! Item {item['id']} has invalid conversation format!")
|
|
return False
|
|
|
|
expected_from = "human" if i % 2 == 0 else "gpt"
|
|
if conv["from"] != expected_from:
|
|
print(f"WARNING!!! Item {item['id']} has incorrect conversation order!")
|
|
return False
|
|
|
|
return True
|
|
|
|
def check_image_and_structure(self, item):
|
|
if not self.check_item_structure(item):
|
|
return
|
|
|
|
# self.check_image_existence(item)
|
|
|
|
def process_images(self):
|
|
if isinstance(self.data, list):
|
|
args = [d for d in self.data]
|
|
with Pool(processes=cpu_count()) as pool:
|
|
list(tqdm(pool.imap(self.check_image_and_structure, args), total=len(self.data)))
|
|
elif isinstance(self.data, dict):
|
|
for d in self.data["datasets"]:
|
|
dd_json_path = d["json_path"]
|
|
data = self.load_json_data(dd_json_path)
|
|
args = [d for d in data]
|
|
with Pool(processes=cpu_count()) as pool:
|
|
list(tqdm(pool.imap(self.check_image_and_structure, args), total=len(data), desc=f"Processing {dd_json_path}"))
|
|
|
|
def count_items(self):
|
|
if isinstance(self.data, list): # Assuming JSON data loaded directly
|
|
return len(self.data)
|
|
elif isinstance(self.data, dict): # Assuming YAML data loaded
|
|
total_items_count = 0
|
|
for d in self.data["datasets"]:
|
|
dd_json_path = d["json_path"]
|
|
data = self.load_json_data(dd_json_path)
|
|
current_items_count = len(data)
|
|
|
|
sampling_strategy = d["sampling_strategy"]
|
|
try:
|
|
if sampling_strategy != "all":
|
|
percentage = float(sampling_strategy.split(":")[-1].replace("%", "")) / 100.0
|
|
else:
|
|
percentage = 1.0
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
percentage = 1.0
|
|
|
|
sampling_count = int(current_items_count * percentage)
|
|
total_items_count += sampling_count
|
|
print(f"{dd_json_path}: {sampling_count}")
|
|
return total_items_count
|
|
|
|
def stat_data(self):
|
|
if isinstance(self.data, dict):
|
|
cur_lens_list = []
|
|
single_image_count = 0
|
|
multiple_image_count = 0
|
|
video_count = 0
|
|
total_count = 0
|
|
text_count = 0
|
|
max_tokens_item = None
|
|
max_tokens = 0
|
|
|
|
for d in self.data["datasets"]:
|
|
dd_json_path = d["json_path"]
|
|
data = self.load_json_data(dd_json_path)
|
|
sampling_strategy = d["sampling_strategy"]
|
|
|
|
try:
|
|
if sampling_strategy != "all":
|
|
percentage = float(sampling_strategy.split(":")[-1].replace("%", "")) / 100.0
|
|
else:
|
|
percentage = 1.0
|
|
except Exception as e:
|
|
print(f"Error parsing sampling strategy: {e}")
|
|
percentage = 1.0
|
|
|
|
sampled_count = int(len(data) * percentage)
|
|
print(f"{dd_json_path}: {sampled_count} (sampled from {len(data)})")
|
|
|
|
for item in data[:sampled_count]:
|
|
conversations = item["conversations"]
|
|
cur_len = sum([len(conv["value"].split()) for conv in conversations])
|
|
cur_lens_list.append(cur_len)
|
|
|
|
if cur_len > max_tokens:
|
|
max_tokens = cur_len
|
|
max_tokens_item = item
|
|
|
|
total_count += 1
|
|
if "image" in item:
|
|
if isinstance(item["image"], list):
|
|
if len(item["image"]) > 1:
|
|
multiple_image_count += 1
|
|
else:
|
|
single_image_count += 1
|
|
else:
|
|
single_image_count += 1
|
|
elif "video" in item:
|
|
video_count += 1
|
|
else:
|
|
text_count += 1
|
|
|
|
print(f"Max length: {max(cur_lens_list)}, Min length: {min(cur_lens_list)}, Average length: {sum(cur_lens_list) / len(cur_lens_list)}")
|
|
print(f"Total items: {total_count}")
|
|
print(f"Text items: {text_count} ({text_count/total_count*100:.2f}%)")
|
|
print(f"Single image items: {single_image_count} ({single_image_count/total_count*100:.2f}%)")
|
|
print(f"Multiple image items: {multiple_image_count} ({multiple_image_count/total_count*100:.2f}%)")
|
|
print(f"Video items: {video_count} ({video_count/total_count*100:.2f}%)")
|
|
|
|
print("\nItem with the largest number of tokens:")
|
|
print(f"Token count: {max_tokens}")
|
|
print("Item content:")
|
|
print(json.dumps(max_tokens_item, indent=2))
|
|
|
|
def filter_data(self):
|
|
if isinstance(self.data, dict):
|
|
for d in self.data["datasets"]:
|
|
dd_json_path = d["json_path"]
|
|
print(f"Processing {dd_json_path}")
|
|
data = self.load_json_data(dd_json_path)
|
|
|
|
filtered_data = []
|
|
mismatch_data = []
|
|
mismatch_flag = False
|
|
for item in data:
|
|
try:
|
|
if "image" in item:
|
|
num_image = len(item["image"]) if isinstance(item["image"], list) else 1
|
|
else:
|
|
num_image = 0
|
|
|
|
if "video" in item:
|
|
num_video = len(item["video"]) if isinstance(item["video"], list) else 1
|
|
else:
|
|
num_video = 0
|
|
|
|
num_visuals = num_image + num_video
|
|
conv_text = ""
|
|
for conv in item["conversations"]:
|
|
conv_text += conv["value"]
|
|
|
|
num_img_token_appearance = conv_text.count("<image>")
|
|
if len(conv_text) == 0:
|
|
print(f"Conversation text is empty for {item}")
|
|
|
|
if num_img_token_appearance == num_visuals or num_img_token_appearance < num_visuals and len(conv_text) > 0:
|
|
filtered_data.append(item)
|
|
elif num_img_token_appearance > num_visuals:
|
|
item["num_img_token_appearance"] = num_img_token_appearance
|
|
item["num_visuals"] = num_visuals
|
|
mismatch_data.append(item)
|
|
|
|
if not mismatch_flag:
|
|
print(f"Data mismatch for {item}")
|
|
|
|
mismatch_flag = True
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
print()
|
|
|
|
if mismatch_flag:
|
|
print(f"Data mismatch for {dd_json_path}")
|
|
|
|
if len(filtered_data) < len(data):
|
|
saving_dd_json_path = dd_json_path.replace(".jsonl", f"fltd_{len(filtered_data)}.json").replace(".json", f"fltd_{len(filtered_data)}.json")
|
|
with open(saving_dd_json_path, "w") as f:
|
|
json.dump(filtered_data, f, indent=2)
|
|
print(f"Filtered data count: {len(filtered_data)}")
|
|
else:
|
|
pass
|
|
|
|
def stat_and_filter_data(self, threshold):
|
|
if isinstance(self.data, dict):
|
|
cur_lens_list = []
|
|
single_image_count = 0
|
|
multiple_image_count = 0
|
|
video_count = 0
|
|
total_count = 0
|
|
text_count = 0
|
|
|
|
for d in self.data["datasets"]:
|
|
dd_json_path = d["json_path"]
|
|
data = self.load_json_data(dd_json_path)
|
|
sampling_strategy = d["sampling_strategy"]
|
|
filtered_data = []
|
|
|
|
try:
|
|
if sampling_strategy != "all":
|
|
percentage = float(sampling_strategy.split(":")[-1].replace("%", "")) / 100.0
|
|
else:
|
|
percentage = 1.0
|
|
except Exception as e:
|
|
print(f"Error parsing sampling strategy: {e}")
|
|
percentage = 1.0
|
|
|
|
sampled_count = int(len(data) * percentage)
|
|
print(f"{dd_json_path}: {sampled_count} (sampled from {len(data)})")
|
|
|
|
save_flag = False
|
|
for item in data:
|
|
total_count += 1
|
|
conversations = item["conversations"]
|
|
filtered_conversations = []
|
|
current_token_count = 0
|
|
|
|
for i in range(0, len(conversations), 2):
|
|
if i + 1 < len(conversations):
|
|
human_conv = conversations[i]
|
|
gpt_conv = conversations[i + 1]
|
|
pair_tokens = len(human_conv["value"].split()) + len(gpt_conv["value"].split())
|
|
|
|
if current_token_count + pair_tokens <= threshold:
|
|
filtered_conversations.extend([human_conv, gpt_conv])
|
|
current_token_count += pair_tokens
|
|
else:
|
|
save_flag = True
|
|
break
|
|
|
|
if filtered_conversations:
|
|
item["conversations"] = filtered_conversations
|
|
cur_len = sum([len(conv["value"].split()) for conv in filtered_conversations])
|
|
cur_lens_list.append(cur_len)
|
|
filtered_data.append(item)
|
|
|
|
if "image" in item:
|
|
if isinstance(item["image"], list):
|
|
if len(item["image"]) > 1:
|
|
multiple_image_count += 1
|
|
else:
|
|
single_image_count += 1
|
|
else:
|
|
single_image_count += 1
|
|
elif "video" in item:
|
|
video_count += 1
|
|
else:
|
|
text_count += 1
|
|
|
|
# Save filtered data for each dataset
|
|
if filtered_data and save_flag:
|
|
if dd_json_path.endswith(".jsonl"):
|
|
output_file = dd_json_path.replace(".jsonl", f"_filtered_{threshold}tokens_{len(filtered_data)}.jsonl")
|
|
with open(output_file, "w") as f:
|
|
for item in filtered_data:
|
|
f.write(json.dumps(item) + "\n")
|
|
else:
|
|
output_file = dd_json_path.replace(".json", f"_filtered_{threshold}tokens_{len(filtered_data)}.json")
|
|
with open(output_file, "w") as f:
|
|
json.dump(filtered_data, f, indent=2)
|
|
print(f"Filtered data for {dd_json_path} saved to: {output_file}")
|
|
|
|
print(f"Max length: {max(cur_lens_list)}, Min length: {min(cur_lens_list)}, Average length: {sum(cur_lens_list) / len(cur_lens_list)}")
|
|
print(f"Total items: {total_count}")
|
|
print(f"Text items: {text_count} ({text_count/total_count*100:.2f}%)")
|
|
print(f"Single image items: {single_image_count} ({single_image_count/total_count*100:.2f}%)")
|
|
print(f"Multiple image items: {multiple_image_count} ({multiple_image_count/total_count*100:.2f}%)")
|
|
print(f"Video items: {video_count} ({video_count/total_count*100:.2f}%)")
|
|
|
|
|
|
def main(file_path, image_root, operation, video_root, threshold=None):
|
|
processor = DataProcessor(file_path, image_root, video_root)
|
|
if operation == "check":
|
|
processor.process_images()
|
|
elif operation == "count":
|
|
total_items = processor.count_items()
|
|
print(f"Total items: {total_items}")
|
|
elif operation == "filter":
|
|
processor.filter_data()
|
|
elif operation == "stat":
|
|
processor.stat_data()
|
|
elif operation == "stat_and_filter":
|
|
if threshold is None:
|
|
raise ValueError("Threshold must be provided for stat_and_filter operation")
|
|
processor.stat_and_filter_data(threshold)
|
|
else:
|
|
raise ValueError("Unsupported operation")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--file_path", type=str, default="/mnt/bn/vl-research/workspace/boli01/projects/LLaVA_Next/scripts/i18n/scale_llms/next_continual.yaml")
|
|
parser.add_argument("--image_root", type=str, default="/mnt/bn/vl-research/data/llava_data")
|
|
parser.add_argument("--video_root", type=str, default="/mnt/bn/vl-research/data/llava_video")
|
|
parser.add_argument("--operation", type=str, default="filter")
|
|
parser.add_argument("--threshold", type=int, default=None, help="Threshold for stat_and_filter operation")
|
|
args = parser.parse_args()
|
|
main(args.file_path, args.image_root, args.operation, args.video_root, args.threshold)
|