Faster indexing

[marcbarb]

I am trying to index large collections more quickly. It seems the doc function in colbert.py

ColBERT/colbert/modeling/colbert.py

Line 94 in 83658dc

def doc(self, input_ids, attention_mask, keep_dims=True):

takes a fair amount of time to run. It looks like the code copies input_ids to the gpu and then the mask function copies them back to the cpu. I've tried rearranging the order, but it doesn't seem to affect overall performance. Is Pytorch caching up the tensors so it doesn't matter where the copy to/from CPU/GPU is done?

Thanks,
Marc Mason

[okhat]

Yes we’re always happy to make things faster. Afaik the time spent in this function is mostly necessary though. It’s primarily spending time on encoding the passages with BERT, not quite something we can optimize away.

you could confirm with profiling that the vast majority of time is spent inside BERT, not in our code.

That said, recent pytorch and transformers optimizations (over the past year or two) are not reflected into this component. You may for instance be able to compile the BERT model into a faster version that’s more static.

I’d look into ONNX BERT.