wehub-resource-sync
144 lines
3.8 KiB
Python
144 lines
3.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
import argparse
|
|
import json
|
|
import os
|
|
|
|
# Third Party
|
|
from tqdm import tqdm
|
|
from transformers import AutoTokenizer
|
|
import numpy as np
|
|
|
|
_tokenizer: AutoTokenizer | None = None
|
|
|
|
|
|
def estimate_num_tokens(text: str) -> int:
|
|
global _tokenizer
|
|
if _tokenizer is None:
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
_tokenizer = AutoTokenizer.from_pretrained(args.model)
|
|
return len(_tokenizer.tokenize(text))
|
|
|
|
|
|
def is_human(human):
|
|
return human in ["human", "user"]
|
|
|
|
|
|
def is_gpt(gpt):
|
|
return gpt in ["gpt", "chatgpt", "bing", "bard"]
|
|
|
|
|
|
def is_system(system):
|
|
return system in ["system"]
|
|
|
|
|
|
def invalid_conversations(conversations):
|
|
# pair does not match
|
|
if len(conversations) < 2:
|
|
return True
|
|
|
|
# starting from gpt or systems
|
|
entry = conversations[0]
|
|
if is_gpt(entry["from"]) or is_system(entry["from"]):
|
|
return True
|
|
|
|
# ending with human
|
|
entry = conversations[-1]
|
|
if is_human(entry["from"]):
|
|
return True
|
|
|
|
prev_from = None
|
|
total_tokens = 0
|
|
for conv in conversations:
|
|
_from = conv["from"]
|
|
total_tokens += estimate_num_tokens(conv["value"])
|
|
|
|
# consecutive rounds (gpt followed by gpt, human followed by human, ..)
|
|
if prev_from == _from:
|
|
return True
|
|
|
|
# too long conversations
|
|
if total_tokens > (128 * 1024):
|
|
return True
|
|
|
|
# unknown from
|
|
if not is_human(_from) and not is_gpt(_from) and not is_system(_from):
|
|
return True
|
|
|
|
prev_from = _from
|
|
|
|
return False
|
|
|
|
|
|
parser = argparse.ArgumentParser(description="Process data percentage.")
|
|
parser.add_argument(
|
|
"--parse",
|
|
type=float,
|
|
default=1,
|
|
help="The percentage of data to process (0 to 1). Default is 1 (100%).",
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default="mistralai/Mistral-7B-Instruct-v0.2",
|
|
help="Model for tokenizer. Default is mistralai/Mistral-7B-Instruct-v0.2.",
|
|
)
|
|
parser.add_argument(
|
|
"--trace",
|
|
type=str,
|
|
default="ShareGPT_V3_unfiltered_cleaned_split.json",
|
|
help="Trace file. Default is ShareGPT_V3_unfiltered_cleaned_split.json",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
print("Loading trace file..")
|
|
with open(args.trace, "r", encoding="utf-8") as file:
|
|
data = json.load(file)
|
|
|
|
num_of_ids = len(data)
|
|
print(f"Number of IDs: {num_of_ids}")
|
|
|
|
# exclude invalid data
|
|
print("Veryfing trace..")
|
|
data = [d for d in tqdm(data) if not invalid_conversations(d["conversations"])]
|
|
excluded_ids = len(data)
|
|
num_of_ids -= excluded_ids
|
|
print(f"Excluded number of IDs: {excluded_ids}")
|
|
|
|
data_to_process = int(num_of_ids * args.parse)
|
|
data = data[:data_to_process]
|
|
print(f"Data to process: {data_to_process}")
|
|
|
|
for d in tqdm(data):
|
|
conversations = d["conversations"]
|
|
d["num_round"] = len(conversations) # human is one round, gpt is another round
|
|
human_tokens = []
|
|
gpt_tokens = []
|
|
for conv in conversations:
|
|
num_tokens = estimate_num_tokens(conv["value"])
|
|
|
|
if is_human(conv["from"]):
|
|
human_tokens.append(num_tokens)
|
|
elif is_gpt(conv["from"]):
|
|
conv["num_tokens"] = num_tokens
|
|
gpt_tokens.append(num_tokens)
|
|
else:
|
|
print("Invalid _from_")
|
|
|
|
if len(human_tokens) == 0:
|
|
d["average_human_token"] = 0
|
|
d["max_human_token"] = 0
|
|
else:
|
|
d["average_human_token"] = float(np.mean(human_tokens))
|
|
d["max_human_token"] = float(np.max(human_tokens))
|
|
if len(gpt_tokens) == 0:
|
|
d["average_gpt_token"] = 0
|
|
d["max_gpt_token"] = 0
|
|
else:
|
|
d["average_gpt_token"] = float(np.mean(gpt_tokens))
|
|
d["max_gpt_token"] = float(np.max(gpt_tokens))
|
|
|
|
with open("ShareGPT.json", "w", encoding="utf-8") as file:
|
|
json.dump(data, file, ensure_ascii=False, indent=2)
|