wehub-resource-sync
150 lines
4.6 KiB
Python
150 lines
4.6 KiB
Python
import argparse
|
|
import collections
|
|
import json
|
|
import os
|
|
import sys
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from glob import glob
|
|
from functools import partial
|
|
|
|
from tqdm import tqdm
|
|
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
|
|
|
|
from data.mathscale.util import mathscale_is_equiv
|
|
|
|
|
|
def majority_voting_frequency(preds):
|
|
assert isinstance(preds, list), preds
|
|
if isinstance(preds[0], list):
|
|
tmp = []
|
|
for pred in preds:
|
|
tmp.append(str(sorted(pred)))
|
|
tmp = collections.Counter(tmp)
|
|
tmp = sorted(tmp.items(), key=lambda x: x[1], reverse=True)
|
|
sorted_preds = [(eval(pred), fre) for pred, fre in tmp]
|
|
elif isinstance(preds[0], str):
|
|
tmp = collections.Counter(preds)
|
|
tmp = sorted(tmp.items(), key=lambda x: x[1], reverse=True)
|
|
sorted_preds = [(pred, fre) for pred, fre in tmp]
|
|
else:
|
|
raise ValueError(f"Unknown type {type(preds[0])}")
|
|
return sorted_preds
|
|
|
|
|
|
def worker(item, n: int):
|
|
sc_pred, sc_freq = majority_voting_frequency(item["pred"][:n])[0]
|
|
sc_res = mathscale_is_equiv(sc_pred, item["label"])[0]
|
|
return {
|
|
"sc_res": sc_res,
|
|
"sc_pred": sc_pred,
|
|
"sc_freq": sc_freq,
|
|
"id": item["id"],
|
|
}
|
|
|
|
|
|
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):
|
|
id2data = {}
|
|
for item in data:
|
|
if item["id"] not in id2data:
|
|
id2data[item["id"]] = item
|
|
continue
|
|
|
|
tmp = id2data[item["id"]]
|
|
if isinstance(tmp["response"], str):
|
|
tmp["response"] = [tmp["response"]]
|
|
if not isinstance(tmp["res"], list):
|
|
tmp["res"] = [tmp["res"]]
|
|
if not isinstance(tmp["pred"], list):
|
|
tmp["pred"] = [tmp["pred"]]
|
|
|
|
tmp["response"] = merge_key(tmp["response"], item["response"])
|
|
tmp["res"] = merge_key(tmp["res"], item["res"])
|
|
tmp["pred"] = merge_key(tmp["pred"], item["pred"])
|
|
assert isinstance(tmp["pred"], list), tmp["pred"]
|
|
id2data[item["id"]] = tmp
|
|
|
|
return list(id2data.values())
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_file", type=str)
|
|
parser.add_argument("--output_file", type=str)
|
|
parser.add_argument("--ps", type=str,
|
|
help="The probabilities of top-1 frequency over total samples, split by comma")
|
|
parser.add_argument("--num_workers", type=int, default=8)
|
|
parser.add_argument("--n", type=int, default=128)
|
|
args = parser.parse_args()
|
|
|
|
if os.path.exists(args.input_file):
|
|
data = json.load(open(args.input_file))
|
|
else:
|
|
data = []
|
|
for file in sorted(glob(args.input_file)):
|
|
print(file)
|
|
data += json.load(open(file))
|
|
if len(data) == 0:
|
|
raise ValueError(f"No data found in {args.input_file}")
|
|
|
|
data = merge_seed_sampled_data(data)
|
|
|
|
pbar = tqdm(data)
|
|
outputs = []
|
|
sc = 0
|
|
num_pred = 0
|
|
freq_pos = collections.Counter()
|
|
freq = collections.Counter()
|
|
with ThreadPoolExecutor(max_workers=args.num_workers) as executor:
|
|
futures = []
|
|
_annotate = partial(worker, n=args.n)
|
|
for _input in pbar:
|
|
future = executor.submit(_annotate, _input)
|
|
num_pred += len(_input["response"])
|
|
futures.append(future)
|
|
pbar.update()
|
|
|
|
for future in tqdm(as_completed(futures), total=len(futures), desc="Collecting results"):
|
|
result = future.result()
|
|
outputs.append(result)
|
|
|
|
if result["sc_res"]:
|
|
sc += 1
|
|
|
|
if result["sc_res"]:
|
|
freq_pos[result["sc_freq"]] += 1
|
|
freq[result["sc_freq"]] += 1
|
|
|
|
ps = list(map(float, args.ps.split(",")))
|
|
reversed_acc_pos = collections.Counter()
|
|
reversed_acc = collections.Counter()
|
|
for p in ps:
|
|
for key, value in freq_pos.items():
|
|
if key / args.n >= p:
|
|
reversed_acc_pos[p] += value
|
|
for key, value in freq.items():
|
|
if key / args.n >= p:
|
|
reversed_acc[p] += value
|
|
|
|
p_ratio = {p: reversed_acc_pos[p] / reversed_acc[p] for p in ps}
|
|
|
|
print(f"Reversed Acc Pos Ratio: {p_ratio}")
|
|
|
|
print(f"SC: {sc}/{len(data)} = {sc / len(data)}")
|
|
print(f"Num Pred: {num_pred}/{len(data)} = {num_pred / len(data)}")
|
|
|
|
json.dump(outputs, open(args.output_file, "w"), ensure_ascii=False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|