chore: import upstream snapshot with attribution

2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
@@ -0,0 +1,145 @@
+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()