def main():
random.seed(12345)
parser = Arguments(description='Faiss indexing for end-to-end retrieval with ColBERT.')
parser.add_index_use_input()
parser.add_argument('--sample', dest='sample', default=None, type=float)
parser.add_argument('--slices', dest='slices', default=1, type=int)
args = parser.parse()
assert args.slices >= 1
assert args.sample is None or (0.0 < args.sample < 1.0), args.sample
with Run.context():
args.index_path = os.path.join(args.index_root, args.index_name)
assert os.path.exists(args.index_path), args.index_path
num_embeddings = sum(load_doclens(args.index_path))
print("#> num_embeddings =", num_embeddings)
if args.partitions is None:
args.partitions = 1 << math.ceil(math.log2(8 * math.sqrt(num_embeddings)))
print('\n\n')
Run.warn("You did not specify --partitions!")
Run.warn("Default computation chooses", args.partitions,
"partitions (for {} embeddings)".format(num_embeddings))
print('\n\n')
index_faiss(args)
def main():
parser = Arguments(
description=
'Training ColBERT with <query, positive passage, negative passage> triples.'
)
parser.add_model_parameters()
parser.add_model_training_parameters()
parser.add_training_input()
args = parser.parse()
assert args.bsize % args.accumsteps == 0, ((
args.bsize, args.accumsteps
), "The batch size must be divisible by the number of gradient accumulation steps."
)
assert args.query_maxlen <= 512
assert args.doc_maxlen <= 512
args.lazy = args.collection is not None
with Run.context(consider_failed_if_interrupted=False):
train(args)
def main():
random.seed(12345)
parser = Arguments(description='Re-ranking over a ColBERT index')
parser.add_model_parameters()
parser.add_model_inference_parameters()
parser.add_reranking_input()
parser.add_index_use_input()
parser.add_argument('--step', dest='step', default=1, type=int)
parser.add_argument('--part-range',
dest='part_range',
default=None,
type=str)
parser.add_argument('--log-scores',
dest='log_scores',
default=False,
action='store_true')
parser.add_argument('--batch',
dest='batch',
default=False,
action='store_true')
parser.add_argument('--depth', dest='depth', default=1000, type=int)
args = parser.parse()
if args.part_range:
part_offset, part_endpos = map(int, args.part_range.split('..'))
args.part_range = range(part_offset, part_endpos)
with Run.context():
args.colbert, args.checkpoint = load_colbert(args)
args.queries = load_queries(args.queries)
args.qrels = load_qrels(args.qrels)
args.topK_pids, args.qrels = load_topK_pids(args.topK,
qrels=args.qrels)
args.index_path = os.path.join(args.index_root, args.index_name)
if args.batch:
batch_rerank(args)
else:
rerank(args)
def main():
random.seed(12345)
parser = Arguments(
description='End-to-end retrieval and ranking with ColBERT.')
parser.add_model_parameters()
parser.add_model_inference_parameters()
parser.add_ranking_input()
parser.add_retrieval_input()
parser.add_argument('--faiss_name',
dest='faiss_name',
default=None,
type=str)
parser.add_argument('--faiss_depth',
dest='faiss_depth',
default=1024,
type=int)
parser.add_argument('--part-range',
dest='part_range',
default=None,
type=str)
parser.add_argument('--batch',
dest='batch',
default=False,
action='store_true')
parser.add_argument('--depth', dest='depth', default=1000, type=int)
args = parser.parse()
args.depth = args.depth if args.depth > 0 else None
if args.part_range:
part_offset, part_endpos = map(int, args.part_range.split('..'))
args.part_range = range(part_offset, part_endpos)
with Run.context():
args.colbert, args.checkpoint = load_colbert(args)
args.qrels = load_qrels(args.qrels)
args.queries = load_queries(args.queries)
args.index_path = os.path.join(args.index_root, args.index_name)
if args.faiss_name is not None:
args.faiss_index_path = os.path.join(args.index_path,
args.faiss_name)
else:
args.faiss_index_path = os.path.join(args.index_path,
get_faiss_index_name(args))
if args.batch:
batch_retrieve(args)
else:
retrieve(args)
def main():
random.seed(12345)
parser = Arguments(
description='Precomputing document representations with ColBERT.')
parser.add_model_parameters()
parser.add_model_inference_parameters()
parser.add_indexing_input()
parser.add_argument('--chunksize',
dest='chunksize',
default=6.0,
required=False,
type=float) # in GiBs
args = parser.parse()
with Run.context():
args.index_path = os.path.join(args.index_root, args.index_name)
# try:
assert not os.path.exists(args.index_path), args.index_path
# except:
# print("\n\nNOT EXISTING:", args.index_path, args.index_path, '\n\n')
distributed.barrier(args.rank)
if args.rank < 1:
create_directory(args.index_root)
create_directory(args.index_path)
distributed.barrier(args.rank)
process_idx = max(0, args.rank)
encoder = CollectionEncoder(args,
process_idx=process_idx,
num_processes=args.nranks)
encoder.encode()
distributed.barrier(args.rank)
# Save metadata.
if args.rank < 1:
metadata_path = os.path.join(args.index_path, 'metadata.json')
print_message("Saving (the following) metadata to", metadata_path,
"..")
print(args.input_arguments)
with open(metadata_path, 'w') as output_metadata:
ujson.dump(args.input_arguments.__dict__, output_metadata)
distributed.barrier(args.rank)
def main():
random.seed(12345)
parser = Arguments(description='Exhaustive (slow, not index-based) evaluation of re-ranking with ColBERT.')
parser.add_model_parameters()
parser.add_model_inference_parameters()
parser.add_reranking_input()
parser.add_argument('--depth', dest='depth', required=False, default=None, type=int)
args = parser.parse()
with Run.context():
args.colbert, args.checkpoint = load_colbert(args)
args.qrels = load_qrels(args.qrels)
if args.collection or args.queries:
assert args.collection and args.queries
args.queries = load_queries(args.queries)
args.collection = load_collection(args.collection)
args.topK_pids, args.qrels = load_topK_pids(args.topK, args.qrels)
else:
args.queries, args.topK_docs, args.topK_pids = load_topK(args.topK)
assert (not args.shortcircuit) or args.qrels, \
"Short-circuiting (i.e., applying minimal computation to queries with no positives in the re-ranked set) " \
"can only be applied if qrels is provided."
evaluate_recall(args.qrels, args.queries, args.topK_pids)
evaluate(args)