Cannot run ColBERT on CPU

[timbmg]

in pyserini/encode/_tct_colbert.py autocast is only imported if CUDA is availabe.

if torch.cuda.is_available():
    from torch.cuda.amp import autocast

However, later is it used in code that should also run on CPU, i.e. here:

pyserini/pyserini/encode/_tct_colbert.py

Line 63 in 760c22a

with autocast():

[lintool]

hi @jacklin64 @justram can you take a look at this please? thanks!

[justram]

Hello @timbmg,

Thank you for bringing this to our attention.
I understand the concern you raised regarding the usage of torch.autocast in the code that should also run on CPU.
You are correct; however, currently, torch.autocast only supports bfloat16 when device='cpu', and our prebuilt index is based on the CUDA + float16 setting.
As a result, encoding the document corpus with device='cpu' & bfloat16 would deviate from our original setting and might lead to significant processing time.
To address this issue and enable CPU-based encoding, we recommend considering the ONNX setting instead.

If you have any further questions or suggestions, please feel free to share them with us.

[zeuner]

A similar situation might be the case for pyserini/encode/_unicoil.py.

Is there any documentation available for encoding (and possibly subsequently indexing) using the ONNX setting to enable CPU support?

[lintool]

We have CPU support in Anserini via ONNX, e.g., https://github.com/castorini/anserini/blob/master/docs/regressions/regressions-msmarco-passage-splade-pp-ed-onnx.md

[lintool]

@timbmg Seeing no further updates, closing. Please re-open if needed.