def setup_task(cls, args, **kwargs):
"""Setup the task (e.g., load dictionaries).
Args:
args (argparse.Namespace): parsed command-line arguments
"""
paths = utils.split_paths(args['task']['data'])
assert len(paths) > 0
# load dictionaries
# src_dict = cls.load_dictionary(os.path.join(paths[0], 'csnjs_8k_9995p_unigram_url.dict.txt'))
src_dict = Dictionary(extra_special_symbols=[
constants.CLS, constants.SEP, constants.MASK, constants.EOL,
constants.URL
])
src_dict.add_from_file(args['dataset']['srcdict'])
tgt_dict = Dictionary.load(args['dataset']['tgtdict'])
# src_dict = cls.load_dictionary(os.path.join(paths[0], '{}.dict.txt'.format(args['task']['source_lang'])))
# tgt_dict = cls.load_dictionary(os.path.join(paths[0], '{}.dict.txt'.format(args['task']['target_lang'])))
# assert src_dict.pad() == tgt_dict.pad()
# assert src_dict.eos() == tgt_dict.eos()
# assert src_dict.unk() == tgt_dict.unk()
# LOGGER.info('[{}] dictionary: {} types'.format(args['task']['source_lang'], len(src_dict)))
# LOGGER.info('[{}] dictionary: {} types'.format(args['task']['target_lang'], len(tgt_dict)))
return cls(args, src_dict, tgt_dict)
def setup_task(cls, args, **kwargs):
"""Setup the task.
"""
# paths = args.data.split(':')
paths = utils.split_paths(args['task']['data'])
assert len(paths) > 0
dictionary = Dictionary.load(os.path.join(paths[0], 'dict.jsonl'))
data_path = paths[0]
if args['task']['langs'] is None:
languages = sorted([
name for name in os.listdir(data_path)
if os.path.isdir(os.path.join(data_path, name))
])
else:
languages = args['task']['langs'] # .split(',')
if args['task']['add_lang_token']:
for lang in languages:
dictionary.add_symbol('[{}]'.format(lang))
LOGGER.info("Loading dictionary: {} types".format(len(dictionary)))
# if not hasattr(args, 'shuffle_instance'):
# args.shuffle_instance = False
return cls(args, dictionary)
def load(cls, f):
subtoken_dict = Dictionary.load(f)
splitted_filenames = f.rsplit('.', 2)
bpe_f = '.'.join([splitted_filenames[0], 'bpe'] +
splitted_filenames[-2:])
bpetoken_dict = WordBpeDicionary.load(bpe_f)
return cls(subtoken_dict, bpetoken_dict)
def save_lang_dict():
src_file = PathManager.expanduser("~/clcdsa/astbert/data-mmap/lang.jsonl")
dict = Dictionary.load(src_file)
tgt_file = os.path.join(args['preprocess']['destdir'], 'lang.jsonl')
PathManager.mkdir(os.path.dirname(tgt_file))
dict.save(tgt_file)
return dict
def load_dictionary(cls, filename):
"""Load the dictionary from the filename
Args:
filename (str): the filename
"""
return Dictionary.load(filename)
def setup_task(cls, args, **kwargs):
"""Setup the task.
"""
dictionary = Dictionary.load(os.path.join(args.data, 'dict.txt'))
LOGGER.info('dictionary: {} types'.format(len(dictionary)))
if not hasattr(args, 'shuffle_instance'):
args.shuffle_instance = False
return cls(args, dictionary)
def load_dictionary(cls, filename):
"""Load the dictionary from the filename
Args:
filename (str): the filename
"""
if filename.endswith('.txt'):
return Dictionary.load(filename)
else:
return Dictionary.load_json(filename)
def load_dictionary(cls, filename):
"""Load the dictionary from the filename
Args:
filename (str): the filename
"""
if filename.endswith('.txt'):
return Dictionary.load(filename)
else:
is_bpe = os.path.basename(filename).split('.')[-3] == 'bpe'
if is_bpe:
return RetrievalDictionary.load_json(filename)
else:
return Dictionary.load_json(filename)
def setup_task(cls, args, **kwargs):
paths = utils.split_paths(args['task']['data'])
assert len(paths) > 0
dictionary = Dictionary.load(os.path.join(paths[0], 'dict.txt'))
LOGGER.info('dictionary: {} types'.format(len(dictionary)))
return cls(args, dictionary)
def cli_main():
SEED = 204
BATCH_SIZE = 64
MAX_SOURCE_POSITIONS = 1024
EPOCH = 50
from ncc.utils.set_seed import set_seed
set_seed(SEED)
use_cuda = torch.cuda.is_available()
if use_cuda:
device = os.environ.get('CUDA_VISIBALE_DEVICES', [0])[0] # get first device as default
torch.cuda.set_device(f'cuda:{device}')
criterion = DeepTuneLoss(task=None, sentence_avg=-1)
if use_cuda:
criterion = criterion.cuda()
data = []
for i, platform in enumerate(LANGUAGES):
DATA_DIR = os.path.join(DATASET_DIR, f'mapping/{platform}/data-mmap')
def get_attr(attr):
oracle_file = os.path.join(DATA_DIR, f'train.{attr}')
with open(oracle_file, 'rb') as reader:
out = pickle.load(reader)
return np.asarray(out)
platform_name = mapping_metrics.platform2str(platform)
benchmarks = get_attr('benchmark')
runtime_cpus = get_attr('runtime_cpu')
runtime_gpus = get_attr('runtime_gpu')
#################### load dataset ####################
src_dataset = load_mmap_dataset(os.path.join(DATA_DIR, f'train.src_tokens'))
src_dataset = TruncateDataset(src_dataset, truncation_length=MAX_SOURCE_POSITIONS, truncate_prefix=0)
tgt_dataset = load_mmap_dataset(os.path.join(DATA_DIR, f'train.oracle'))
src_dict = Dictionary.load(os.path.join(DATA_DIR, 'src_tokens.dict.jsonl'))
src_aux = OrderedDict()
src_aux['transfer'] = get_attr('transfer')
src_aux['wgsize'] = get_attr('wgsize')
tgt_dict = Dictionary.load(os.path.join(DATA_DIR, 'oracle.dict.jsonl'))
dataset = LanguagePairDataset(
src=src_dataset, src_sizes=src_dataset.sizes, src_dict=src_dict, src_aux=src_aux,
tgt=tgt_dataset, tgt_sizes=tgt_dataset.sizes, tgt_dict=tgt_dict, tgt_aux=None,
left_pad_source=True, max_source_positions=MAX_SOURCE_POSITIONS,
)
#################### load dataset ####################
# build toy dataset for 10-fold cross validation
tgt_data = [tgt_dataset[idx].item() for idx in range(len(tgt_dataset))]
src_data = [None] * len(tgt_data)
# 10-fold cross-validation
kf = StratifiedKFold(n_splits=10, shuffle=True, random_state=SEED)
for j, (train_ids, test_ids) in enumerate(kf.split(src_data, tgt_data)):
# deeptune model
model = DeepTuneEncoder(dictionary=src_dict, embed_dim=64,
rnn_cell='lstm', rnn_hidden_dim=64, rnn_dropout=0., rnn_num_layers=2,
aux_dim=2, inner_dim=32, out_dim=2)
if use_cuda:
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch_i in range(EPOCH):
if dataset.shuffle:
random.shuffle(train_ids)
train_batch_sampler = data_utils.batch_by_size(
train_ids,
num_tokens_fn=lambda *args: -1,
max_sentences=BATCH_SIZE,
)
train_dataloader = DataLoader(dataset=dataset,
batch_sampler=train_batch_sampler,
collate_fn=collate, )
with tqdm(total=len(train_dataloader)) as t:
for sample_i, sample in enumerate(train_dataloader, start=1):
t.set_description(f'Epoch {epoch_i + 1}/{EPOCH} Batch {sample_i}/{len(train_dataloader)}')
if use_cuda:
sample = move_to_cuda(sample)
loss, sample_size, logging_output = criterion(model, sample)
loss.div_(sample_size)
t.set_postfix(loss=loss.item())
t.update()
optimizer.zero_grad()
loss.backward()
optimizer.step()
# test accuracy
test_batch_sampler = data_utils.batch_by_size(
test_ids,
num_tokens_fn=lambda *args: -1,
max_sentences=BATCH_SIZE,
)
test_dataloader = DataLoader(dataset=dataset,
batch_sampler=test_batch_sampler,
collate_fn=collate, )
predictions, ground_truth = [], []
for sample in test_dataloader:
if use_cuda:
sample = move_to_cuda(sample)
hybrid_out, _ = model(**sample['net_input'])
predictions.append(hybrid_out.max(dim=-1)[1])
ground_truth.append(sample['target'].view(-1))
predictions = torch.cat(predictions)
ground_truth = torch.cat(ground_truth)
accuracy = (predictions == ground_truth).tolist()
# runtimes of baseline mapping (CPU on AMD, GPU on NVIDIA)
gt_runtimes = (runtime_cpus if platform == "amd" else runtime_gpus)[test_ids]
pred_runtimes = [
(runtime_cpus if pred == 0 else runtime_gpus)[idx]
for idx, pred in zip(test_ids, predictions)
]
speedup = gt_runtimes / pred_runtimes
# record results
for benchmark_, o_, p_, accuracy_, p_speedup_ in \
zip(benchmarks[test_ids], ground_truth, predictions, accuracy, speedup):
data.append({
"Model": model.__class__.__name__,
"Platform": platform_name,
'Benchmark': mapping_metrics.escape_benchmark_name(benchmark_),
'Benchmark Suite': mapping_metrics.escape_suite_name(benchmark_),
"Oracle Mapping": o_,
"Predicted Mapping": p_,
"Accuracy": accuracy_,
"Speedup": p_speedup_,
})
del model, optimizer
performance = pd.DataFrame(
data, index=range(1, len(data) + 1), columns=[
"Model",
"Platform",
"Benchmark",
"Benchmark Suite",
"Oracle Mapping",
"Predicted Mapping",
"Accuracy",
"Speedup"
])
benchmark_out = performance.groupby(['Platform', 'Benchmark Suite'])[['Platform', 'Accuracy', 'Speedup']].mean()
benchmark_out['Accuracy'] = round(benchmark_out['Accuracy'] * 100, 2)
benchmark_out['Speedup'] = round(benchmark_out['Speedup'], 2)
print(benchmark_out)
out = performance.groupby(['Platform'])[['Platform', 'Accuracy', 'Speedup']].mean()
out['Accuracy'] = round(out['Accuracy'] * 100, 2)
out['Speedup'] = round(out['Speedup'], 2)
print(out)
def main(args):
LOGGER.info('mkdir {} for {} task'.format(args['preprocess']['destdir'],
args['preprocess']['task']))
PathManager.mkdir(args['preprocess']['destdir'])
from ncc.data.dictionary import TransformersDictionary
vocab = TransformersDictionary.from_pretrained(
'microsoft/graphcodebert-base')
file = os.path.join(args['preprocess']['destdir'], 'dfs.jsonl')
node_dict = Dictionary.load(file)
# 2. ***************build dataset********************
# dump into pkl file
# transform a language's code into src format and tgt format simualtaneouly
num_workers = args['preprocess']['workers']
src_lang, tgt_lang = args['preprocess']['src_lang'], args['preprocess'][
'tgt_lang']
# code tokens => code tokens
for mode, lang in itertools.product(MODES, [src_lang, tgt_lang]):
data_dir = str.replace(args['preprocess'][f'{mode}pref'], '*', lang)
src_file = f"{data_dir}.code_tokens"
PathManager.mkdir(os.path.join(args['preprocess']['destdir'], lang))
dst_file = os.path.join(args['preprocess']['destdir'], lang,
f"{mode}.code_tokens")
offsets = find_offsets(src_file, num_workers)
pool = None
if num_workers > 1:
# p1-pN -> (1 bin-txt, 1 idx), (N bin-txt, N idx)
pool = Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(dst_file, worker_id)
pool.apply_async(
binarize_tokens,
(args, src_file, vocab, prefix, offsets[worker_id],
offsets[worker_id + 1]),
)
pool.close()
ds_file = '{}.mmap'.format(dst_file)
ds = indexed_dataset.make_builder(ds_file,
impl="mmap",
vocab_size=len(vocab))
end = offsets[1]
with file_io.open(src_file, 'r') as reader:
reader.seek(0)
line = file_io.safe_readline(reader)
while line:
if end > 0 and reader.tell() > end:
break
line = json_io.json_loads(line)
code_tokens = vocab.subtokenize(line)
code_tokens = torch.IntTensor(
vocab.tokens_to_indices(code_tokens))
ds.add_item(code_tokens)
line = reader.readline()
if num_workers > 1:
# p1-pN
pool.join()
# merge sub-processors' index and data files into final files and delete them.
for worker_id in range(1, num_workers):
temp_file_path = "{}{}".format(dst_file, worker_id)
ds.merge_file_(temp_file_path)
# idx, txt
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize('{}.idx'.format(dst_file))
# code => code
for mode, lang in itertools.product(MODES, [src_lang, tgt_lang]):
data_dir = str.replace(args['preprocess'][f'{mode}pref'], '*', lang)
src_file = f"{data_dir}.code"
PathManager.mkdir(os.path.join(args['preprocess']['destdir'], lang))
dst_file = os.path.join(args['preprocess']['destdir'], lang,
f"{mode}.code")
ds_file = '{}.bin'.format(dst_file)
ds = indexed_dataset.make_builder(ds_file,
impl="bin",
vocab_size=len(vocab))
with open(src_file, 'r') as reader:
for line in reader:
line = json_io.json_loads(line)
ds.add_item(line)
ds.finalize('{}.idx'.format(dst_file))
# dfs => dfs
for mode, lang in itertools.product(MODES, [src_lang, tgt_lang]):
data_dir = str.replace(args['preprocess'][f'{mode}pref'], '*', lang)
src_file = f"{data_dir}.dfs"
PathManager.mkdir(os.path.join(args['preprocess']['destdir'], lang))
dst_file = os.path.join(args['preprocess']['destdir'], lang,
f"{mode}.dfs")
offsets = find_offsets(src_file, num_workers)
pool = None
if num_workers > 1:
# p1-pN -> (1 bin-txt, 1 idx), (N bin-txt, N idx)
pool = Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(dst_file, worker_id)
pool.apply_async(
binarize_dfs,
(args, src_file, node_dict, prefix, offsets[worker_id],
offsets[worker_id + 1]),
)
pool.close()
ds_file = '{}.mmap'.format(dst_file)
ds = indexed_dataset.make_builder(ds_file,
impl="mmap",
vocab_size=len(vocab))
end = offsets[1]
with file_io.open(src_file, 'r') as reader:
reader.seek(0)
line = file_io.safe_readline(reader)
while line:
if end > 0 and reader.tell() > end:
break
line = json_io.json_loads(line)
dfs = torch.IntTensor([node_dict.index(tok) for tok in line])
ds.add_item(dfs)
line = reader.readline()
if num_workers > 1:
# p1-pN
pool.join()
# merge sub-processors' index and data files into final files and delete them.
for worker_id in range(1, num_workers):
temp_file_path = "{}{}".format(dst_file, worker_id)
ds.merge_file_(temp_file_path)
# idx, txt
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize('{}.idx'.format(dst_file))
