39 lines
1.7 KiB
Python
39 lines
1.7 KiB
Python
import json
|
|
import os
|
|
from itertools import islice
|
|
import heapq
|
|
|
|
from tqdm import tqdm
|
|
|
|
if __name__ == '__main__':
|
|
mal_file_name = '/home/king/python/data/jsonl/infected_jsonl/'
|
|
ben_file_name = '/home/king/python/data/jsonl/refind_jsonl/'
|
|
fun_name_dict = {}
|
|
for file in tqdm(os.listdir(mal_file_name)):
|
|
with open(mal_file_name + file, 'r') as item:
|
|
item = json.loads(item.readline())
|
|
item_fun_list = item['function_names']
|
|
for fun_name in item_fun_list:
|
|
if fun_name != 'start' and fun_name != 'start_0' and 'sub_' not in fun_name:
|
|
if fun_name_dict.get(fun_name) is not None:
|
|
fun_name_dict[fun_name] += 1
|
|
else:
|
|
fun_name_dict[fun_name] = 1
|
|
for file in tqdm(os.listdir(ben_file_name)):
|
|
with open(ben_file_name + file, 'r') as item:
|
|
item = json.loads(item.readline())
|
|
item_fun_list = item['function_names']
|
|
for fun_name in item_fun_list:
|
|
if fun_name != 'start' and fun_name != 'start_0' and 'sub_' not in fun_name:
|
|
if fun_name_dict.get(fun_name) is not None:
|
|
fun_name_dict[fun_name] += 1
|
|
else:
|
|
fun_name_dict[fun_name] = 1
|
|
with open('/home/king/python/data/processed_dataset/train_external_function_name_vocab.jsonl', 'w') as file:
|
|
largest_10000_items = heapq.nlargest(10000, fun_name_dict.items(), key=lambda item: item[1])
|
|
for key, value in largest_10000_items:
|
|
temp = {"f_name": key, "count": value}
|
|
file.write(json.dumps(temp) + '\n')
|
|
|
|
|