|
import os
|
|
import tqdm
|
|
import torch
|
|
import datetime
|
|
import itertools
|
|
|
|
from multiprocessing import Pool
|
|
from collections import OrderedDict, defaultdict
|
|
|
|
|
|
def print_message(*s, condition=True):
|
|
s = ' '.join([str(x) for x in s])
|
|
msg = "[{}] {}".format(datetime.datetime.now().strftime("%b %d, %H:%M:%S"), s)
|
|
|
|
if condition:
|
|
print(msg, flush=True)
|
|
|
|
return msg
|
|
|
|
|
|
def timestamp():
|
|
format_str = "%Y-%m-%d_%H.%M.%S"
|
|
result = datetime.datetime.now().strftime(format_str)
|
|
return result
|
|
|
|
|
|
def file_tqdm(file):
|
|
print(f"#> Reading {file.name}")
|
|
|
|
with tqdm.tqdm(total=os.path.getsize(file.name) / 1024.0 / 1024.0, unit="MiB") as pbar:
|
|
for line in file:
|
|
yield line
|
|
pbar.update(len(line) / 1024.0 / 1024.0)
|
|
|
|
pbar.close()
|
|
|
|
|
|
def save_checkpoint(path, epoch_idx, mb_idx, model, optimizer, arguments=None):
|
|
print(f"#> Saving a checkpoint to {path} ..")
|
|
|
|
if hasattr(model, 'module'):
|
|
model = model.module
|
|
|
|
checkpoint = {}
|
|
checkpoint['epoch'] = epoch_idx
|
|
checkpoint['batch'] = mb_idx
|
|
checkpoint['model_state_dict'] = model.state_dict()
|
|
checkpoint['optimizer_state_dict'] = optimizer.state_dict()
|
|
checkpoint['arguments'] = arguments
|
|
|
|
torch.save(checkpoint, path)
|
|
|
|
|
|
def load_checkpoint(path, model, optimizer=None, do_print=True):
|
|
if do_print:
|
|
print_message("#> Loading checkpoint", path, "..")
|
|
|
|
if path.startswith("http:") or path.startswith("https:"):
|
|
checkpoint = torch.hub.load_state_dict_from_url(path, map_location='cpu')
|
|
else:
|
|
checkpoint = torch.load(path, map_location='cpu')
|
|
|
|
state_dict = checkpoint['model_state_dict']
|
|
new_state_dict = OrderedDict()
|
|
for k, v in state_dict.items():
|
|
name = k
|
|
if k[:7] == 'module.':
|
|
name = k[7:]
|
|
new_state_dict[name] = v
|
|
|
|
checkpoint['model_state_dict'] = new_state_dict
|
|
|
|
try:
|
|
model.load_state_dict(checkpoint['model_state_dict'])
|
|
except:
|
|
print_message("[WARNING] Loading checkpoint with strict=False")
|
|
model.load_state_dict(checkpoint['model_state_dict'], strict=False)
|
|
|
|
if optimizer:
|
|
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
|
|
|
if do_print:
|
|
print_message("#> checkpoint['epoch'] =", checkpoint['epoch'])
|
|
print_message("#> checkpoint['batch'] =", checkpoint['batch'])
|
|
|
|
return checkpoint
|
|
|
|
|
|
def create_directory(path):
|
|
if os.path.exists(path):
|
|
print('\n')
|
|
print_message("#> Note: Output directory", path, 'already exists\n\n')
|
|
else:
|
|
print('\n')
|
|
print_message("#> Creating directory", path, '\n\n')
|
|
os.makedirs(path)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def f7(seq):
|
|
"""
|
|
Source: https://stackoverflow.com/a/480227/1493011
|
|
"""
|
|
|
|
seen = set()
|
|
return [x for x in seq if not (x in seen or seen.add(x))]
|
|
|
|
|
|
def batch(group, bsize, provide_offset=False):
|
|
offset = 0
|
|
while offset < len(group):
|
|
L = group[offset: offset + bsize]
|
|
yield ((offset, L) if provide_offset else L)
|
|
offset += len(L)
|
|
return
|
|
|
|
|
|
class dotdict(dict):
|
|
"""
|
|
dot.notation access to dictionary attributes
|
|
Credit: derek73 @ https://stackoverflow.com/questions/2352181
|
|
"""
|
|
__getattr__ = dict.__getitem__
|
|
__setattr__ = dict.__setitem__
|
|
__delattr__ = dict.__delitem__
|
|
|
|
|
|
def flatten(L):
|
|
return [x for y in L for x in y]
|
|
|
|
|
|
def zipstar(L, lazy=False):
|
|
"""
|
|
A much faster A, B, C = zip(*[(a, b, c), (a, b, c), ...])
|
|
May return lists or tuples.
|
|
"""
|
|
|
|
if len(L) == 0:
|
|
return L
|
|
|
|
width = len(L[0])
|
|
|
|
if width < 100:
|
|
return [[elem[idx] for elem in L] for idx in range(width)]
|
|
|
|
L = zip(*L)
|
|
|
|
return L if lazy else list(L)
|
|
|
|
|
|
def zip_first(L1, L2):
|
|
length = len(L1) if type(L1) in [tuple, list] else None
|
|
|
|
L3 = list(zip(L1, L2))
|
|
|
|
assert length in [None, len(L3)], "zip_first() failure: length differs!"
|
|
|
|
return L3
|
|
|
|
|
|
def int_or_float(val):
|
|
if '.' in val:
|
|
return float(val)
|
|
|
|
return int(val)
|
|
|
|
def load_ranking(path, types=None, lazy=False):
|
|
print_message(f"#> Loading the ranked lists from {path} ..")
|
|
|
|
try:
|
|
lists = torch.load(path)
|
|
lists = zipstar([l.tolist() for l in tqdm.tqdm(lists)], lazy=lazy)
|
|
except:
|
|
if types is None:
|
|
types = itertools.cycle([int_or_float])
|
|
|
|
with open(path) as f:
|
|
lists = [[typ(x) for typ, x in zip_first(types, line.strip().split('\t'))]
|
|
for line in file_tqdm(f)]
|
|
|
|
return lists
|
|
|
|
|
|
def save_ranking(ranking, path):
|
|
lists = zipstar(ranking)
|
|
lists = [torch.tensor(l) for l in lists]
|
|
|
|
torch.save(lists, path)
|
|
|
|
return lists
|
|
|
|
|
|
def groupby_first_item(lst):
|
|
groups = defaultdict(list)
|
|
|
|
for first, *rest in lst:
|
|
rest = rest[0] if len(rest) == 1 else rest
|
|
groups[first].append(rest)
|
|
|
|
return groups
|
|
|
|
|
|
def process_grouped_by_first_item(lst):
|
|
"""
|
|
Requires items in list to already be grouped by first item.
|
|
"""
|
|
|
|
groups = defaultdict(list)
|
|
|
|
started = False
|
|
last_group = None
|
|
|
|
for first, *rest in lst:
|
|
rest = rest[0] if len(rest) == 1 else rest
|
|
|
|
if started and first != last_group:
|
|
yield (last_group, groups[last_group])
|
|
assert first not in groups, f"{first} seen earlier --- violates precondition."
|
|
|
|
groups[first].append(rest)
|
|
|
|
last_group = first
|
|
started = True
|
|
|
|
return groups
|
|
|
|
|
|
def grouper(iterable, n, fillvalue=None):
|
|
"""
|
|
Collect data into fixed-length chunks or blocks
|
|
Example: grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
|
|
Source: https://docs.python.org/3/library/itertools.html#itertools-recipes
|
|
"""
|
|
|
|
args = [iter(iterable)] * n
|
|
return itertools.zip_longest(*args, fillvalue=fillvalue)
|
|
|
|
|
|
|
|
class NullContextManager(object):
|
|
def __init__(self, dummy_resource=None):
|
|
self.dummy_resource = dummy_resource
|
|
def __enter__(self):
|
|
return self.dummy_resource
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
|
|
def load_batch_backgrounds(args, qids):
|
|
if args.qid2backgrounds is None:
|
|
return None
|
|
|
|
qbackgrounds = []
|
|
|
|
for qid in qids:
|
|
back = args.qid2backgrounds[qid]
|
|
|
|
if len(back) and type(back[0]) == int:
|
|
x = [args.collection[pid] for pid in back]
|
|
else:
|
|
x = [args.collectionX.get(pid, '') for pid in back]
|
|
|
|
x = ' [SEP] '.join(x)
|
|
qbackgrounds.append(x)
|
|
|
|
return qbackgrounds
|
|
|
