|
|
import sys |
|
|
import random |
|
|
import os |
|
|
import re |
|
|
from tqdm import tqdm |
|
|
|
|
|
def filter_poly(asm, eqn): |
|
|
rejects = {"ln", "exp", "sin", "cos", "sqrt", "tan", "asin", "acos", "atan", "E", "pi", "cot"} |
|
|
return any(t in rejects for t in asm.strip().split(" ")) \ |
|
|
or any(t in rejects for t in eqn.strip().split(" ")) |
|
|
|
|
|
def filter_bigint(asm, eqn): |
|
|
if re.search(r"CONST=[0-9]{4,}", asm): |
|
|
return True |
|
|
return False |
|
|
|
|
|
if __name__ == "__main__": |
|
|
import argparse |
|
|
parser = argparse.ArgumentParser("Deduplicate ASM and split files into train/test/valid") |
|
|
parser.add_argument("--inprefix", required=True, help="Prefix of input files") |
|
|
parser.add_argument("--outdir", required=True) |
|
|
parser.add_argument("--split", type=float, default=0.05) |
|
|
parser.add_argument("--seed", type=int, default=1225) |
|
|
parser.add_argument("--filter", choices=["poly", "bigint"], default=None) |
|
|
parser.add_argument("--no-separate-eqn", action="store_true") |
|
|
|
|
|
args = parser.parse_args() |
|
|
|
|
|
eq_mapped = {} |
|
|
combined_ds = [] |
|
|
asm_hash = set() |
|
|
removed = 0 |
|
|
|
|
|
with open(args.inprefix + ".asm", "r") as asmf, \ |
|
|
open(args.inprefix + ".eqn", "r") as eqnf, \ |
|
|
open(args.inprefix + ".const.jsonl", "r") as constf: |
|
|
for i, (asm, eqn, const) in tqdm(enumerate(zip(asmf, eqnf, constf)), |
|
|
desc="Read files", leave=False): |
|
|
h = hash(asm) |
|
|
if h in asm_hash: |
|
|
|
|
|
removed += 1 |
|
|
continue |
|
|
|
|
|
if re.search(r"[0-9]\.[0-9]", eqn): |
|
|
|
|
|
removed += 1 |
|
|
continue |
|
|
|
|
|
if args.filter == "poly" and filter_poly(asm, eqn): |
|
|
removed += 1 |
|
|
continue |
|
|
if args.filter == "bigint" and filter_bigint(asm, eqn): |
|
|
removed += 1 |
|
|
continue |
|
|
|
|
|
asm_hash.add(h) |
|
|
if args.no_separate_eqn: |
|
|
combined_ds.append((i, asm, eqn, const)) |
|
|
else: |
|
|
if eqn not in eq_mapped: |
|
|
eq_mapped[eqn] = [] |
|
|
eq_mapped[eqn].append((i, asm, const)) |
|
|
|
|
|
print("Removed", removed) |
|
|
|
|
|
if args.no_separate_eqn: |
|
|
dataset = combined_ds |
|
|
else: |
|
|
dataset = list(eq_mapped.keys()) |
|
|
|
|
|
random.seed(args.seed) |
|
|
random.shuffle(dataset) |
|
|
|
|
|
N = len(dataset) |
|
|
Ntest = int(N * args.split) |
|
|
|
|
|
splits = { |
|
|
"train": dataset[:N-2*Ntest], |
|
|
"valid": dataset[N-2*Ntest:N-Ntest], |
|
|
"test": dataset[N-Ntest:] |
|
|
} |
|
|
splitidxs = {s: [] for s in splits} |
|
|
|
|
|
idxf = open(os.path.join(args.outdir, "splits.txt"), "w") |
|
|
for s in splits: |
|
|
asmfn = os.path.join(args.outdir, f"{s}.asm") |
|
|
eqnfn = os.path.join(args.outdir, f"{s}.eqn") |
|
|
constfn = os.path.join(args.outdir, f"{s}.const.jsonl") |
|
|
with open(asmfn, "w") as asmf, open(eqnfn, "w") as eqnf, \ |
|
|
open(constfn, "w") as constf: |
|
|
if args.no_separate_eqn: |
|
|
for i, asm, eqn, const in splits[s]: |
|
|
asmf.write(asm) |
|
|
eqnf.write(eqn) |
|
|
constf.write(const) |
|
|
splitidxs[s].append(i) |
|
|
else: |
|
|
for eqn in splits[s]: |
|
|
for i, asm, const in eq_mapped[eqn]: |
|
|
asmf.write(asm) |
|
|
eqnf.write(eqn) |
|
|
constf.write(const) |
|
|
splitidxs[s].append(i) |
|
|
print("Split", s, len(splitidxs[s])) |
|
|
idxf.write(f"==== {s} ====\n") |
|
|
for j, i in enumerate(splitidxs[s]): |
|
|
idxf.write(f"{j}: {i}\n") |
|
|
idxf.write("\n") |
|
|
idxf.close() |
|
|
|
|
|
|
