wehub-resource-sync
146 lines
4.2 KiB
Python
146 lines
4.2 KiB
Python
import json
|
|
from transformers import AutoTokenizer, PreTrainedTokenizer
|
|
import argparse
|
|
from glob import glob
|
|
from tqdm import tqdm
|
|
import os
|
|
from multiprocessing.pool import Pool
|
|
import matplotlib.pyplot as plt
|
|
|
|
_tokenizer: PreTrainedTokenizer
|
|
|
|
|
|
def _init_(tokenizer):
|
|
global _tokenizer
|
|
_tokenizer = tokenizer
|
|
|
|
|
|
def plot_histogram(data, bins=10, x_label="Value", y_label="Frequency", title="Histogram", output_file="histogram.png"):
|
|
# clear previous data
|
|
plt.clf()
|
|
plt.hist(data, bins=bins, edgecolor='black', alpha=0.7)
|
|
plt.xlabel(x_label)
|
|
plt.ylabel(y_label)
|
|
plt.title(title)
|
|
plt.grid(True)
|
|
# plt.show()
|
|
plt.savefig(output_file)
|
|
|
|
|
|
def merge_key(item, value):
|
|
assert isinstance(item, list)
|
|
if isinstance(value, list):
|
|
item = item + value
|
|
else:
|
|
item.append(value)
|
|
return item
|
|
|
|
|
|
def merge_seed_sampled_data(data, key_field="response"):
|
|
id2data = {}
|
|
for item in data:
|
|
if item["id"] not in id2data:
|
|
id2data[item["id"]] = item
|
|
continue
|
|
|
|
tmp = id2data[item["id"]]
|
|
if isinstance(tmp[key_field], str):
|
|
tmp[key_field] = [tmp[key_field]]
|
|
|
|
tmp[key_field] = merge_key(tmp[key_field], item[key_field])
|
|
id2data[item["id"]] = tmp
|
|
|
|
return list(id2data.values())
|
|
|
|
|
|
def worker(item):
|
|
text = item["text"]
|
|
|
|
tokens = _tokenizer.tokenize(text)
|
|
|
|
item["length"] = len(tokens)
|
|
return item
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_file", type=str)
|
|
parser.add_argument("--tokenizer", "-t", type=str)
|
|
parser.add_argument("--key_field", type=str, default="response")
|
|
parser.add_argument("--topic_field", type=str, default=None)
|
|
parser.add_argument("--ks", type=str, default="1,4,8,16")
|
|
parser.add_argument("--num_workers", type=int, default=16)
|
|
# parser.add_argument("--output_file", type=str, default="response_length.png")
|
|
args = parser.parse_args()
|
|
|
|
if os.path.exists(args.input_file):
|
|
print("Reading from file")
|
|
print(args.input_file)
|
|
with open(args.input_file, "r") as f:
|
|
data = json.load(f)
|
|
else:
|
|
data = []
|
|
for file in glob(args.input_file):
|
|
print(file)
|
|
with open(file, "r") as f:
|
|
data.extend(json.load(f))
|
|
|
|
data = merge_seed_sampled_data(data, key_field=args.key_field)
|
|
ks = sorted([int(k) for k in args.ks.split(",")])
|
|
ks = [0] + ks
|
|
mp_inputs = []
|
|
for item in data:
|
|
if isinstance(item[args.key_field], str):
|
|
item[args.key_field] = [item[args.key_field]]
|
|
|
|
_inputs = [{"text": x} for x in item[args.key_field]]
|
|
if args.topic_field:
|
|
for x in _inputs:
|
|
x["topic"] = item[args.topic_field]
|
|
|
|
for i, k in enumerate(ks):
|
|
if i == 0:
|
|
continue
|
|
for x in _inputs[ks[i - 1]:k]:
|
|
x["id"] = item["id"]
|
|
x["k"] = k
|
|
mp_inputs.append(x)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
|
|
with Pool(args.num_workers, initializer=_init_, initargs=(tokenizer,)) as p:
|
|
results = list(tqdm(p.imap(worker, mp_inputs), total=len(mp_inputs)))
|
|
|
|
k2data = {k: [] for k in ks}
|
|
for item in results:
|
|
k2data[item["k"]].append(item["length"])
|
|
|
|
acc = 0
|
|
acc_n = 0
|
|
for k, data in k2data.items():
|
|
acc += sum(data)
|
|
acc_n += len(data)
|
|
if acc_n:
|
|
print(f"k={k}, len={acc_n}, average={acc / acc_n}")
|
|
else:
|
|
print(f"k={k}, len={acc_n}, average=0")
|
|
|
|
if args.topic_field:
|
|
topic2data = {}
|
|
for item in results:
|
|
topic = item["topic"]
|
|
if topic not in topic2data:
|
|
topic2data[topic] = []
|
|
topic2data[topic].append(item["length"])
|
|
|
|
for topic, data in topic2data.items():
|
|
if len(data):
|
|
print(f"topic={topic}, len={len(data)}, average={sum(data) / len(data)}")
|
|
else:
|
|
print(f"topic={topic}, len={len(data)}, average=0")
|
|
|
|
plot_histogram(data, bins=10, x_label="Length", y_label="Frequency", title=f"{topic} Histogram", output_file=f"{topic}_histogram.png")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|