Support CUDA gpus

[irthomasthomas]

Hi Simon,

I was just looking at https://huggingface.co/datasets/SciPhi/AgentSearch-V1
1 Billion passages and 50 Million documents with jina embeddings.

So it motivated me to add cuda support to your jina plugin. CPU is way too slow for me. Only problem is you need to drop the batch size to 1 or 2 for most consumer devices.

Changes

• Add try/except block to move model to CUDA or keep it on CPU if CUDA is not available
• No changes to functionality or API of the class

Timing tests

CPU - Batch size 10 (The maximum I can run in 64GB)

time llm embed-multi prompts-cpu-b10-run1 -d embeddings.db --attach logs $(llm logs path) --sql 'SELECT id, prompt FROM logs.responses LIMIT 1000' -m jina-embeddings-v2-base-en --prefix prompt/ --batch-size 10

Output:

llm embed-multi prompts-cpu-b10-run1 -d embeddings.db --attach logs  --sql  -m
  7142.22s user 3449.01s system 1320% cpu 13:21.98 total

GPU

time llm embed-multi prompts-gpu-b1-run2 -d embeddings.db --attach logs  --sql  -m jina-embeddings-v2-base-en --prefix prompt/ --batch-size 1

Output:

llm embed-multi prompts-gpu-b1-run2 -d embeddings.db --attach logs  --sql  -m  
25.44s user 11.39s system 130% cpu 28.310 total

[irthomasthomas]

Hi @simonw I'd like to add a flag so we can choose cpu or cuda (and does metal need its own flag or will cpu work?) but I'm not familiar with click. I can make simple click apps with chatgpt, but I got lost trying to figure out exactly where and how to add a new click option to your project. Could you point me in the right direction? Thanks!