Updated to allow the selection of GPU for embedding where there is mo… #1734
Conversation
…re than one available. Defaults to cuda[0] or cpu if cuda is not available. Commented reference in settings.yaml under embedding.
| @@ -7,7 +7,7 @@ | |||
| from private_gpt.settings.settings import Settings | |||
|
|
|||
| logger = logging.getLogger(__name__) | |||
|
|
|||
| import torch | |||
There was a problem hiding this comment.
I'd move this to the try block within "huggingface" case. There is no "torch" general dependency declared in pyproject.toml, so this could break the whole execution for people not using huggingface. Actually, we may need to add torch to embeddings-huggingface = ["llama-index-embeddings-huggingface"] as
# Optional Huggingface related dependency
torch = {version = "^2.2.1", optional = true}
embeddings-huggingface = ["torch", "llama-index-embeddings-huggingface"]
in pyproject.toml.
I think huggingface package from llamaindex already depends on torch, but given we are now importing it explicitly we should also depende on it.
| device = torch.device("cuda:0") | ||
| else: | ||
| # If CUDA is not available, use CPU | ||
| device = torch.device("cpu") |
There was a problem hiding this comment.
What happens with laptops using a GPU that is not Nvidia based? For example Mac book running Metal GPU? Will this make embedding slower forcing them to go to CPU?
There was a problem hiding this comment.
This logic looks similar this: llama_index.core.utils.infer_torch_device which handles Metal (mps).
settings.yaml
Outdated
| @@ -54,6 +54,7 @@ embedding: | |||
| # Should be matching the value above in most cases | |||
| mode: huggingface | |||
| ingest_mode: simple | |||
| # gpu: cuda[0] # if you have more than one GPU and you want to select another. defaults to cuda[0], or cpu if cuda not available | |||
There was a problem hiding this comment.
You'd need to include this new setting to settings.py
| @@ -28,9 +28,33 @@ def __init__(self, settings: Settings) -> None: | |||
| "Local dependencies not found, install with `poetry install --extras embeddings-huggingface`" | |||
| ) from e | |||
|
|
|||
| # Get the number of available GPUs | |||
| num_gpus = torch.cuda.device_count() | |||
There was a problem hiding this comment.
Adding code to the codebase just to print information is not a good practive. I'd remove this. whole block of prints.
Updated to allow the selection of GPU for embedding where there is more than one available. Defaults to cuda[0] or cpu if cuda is not available. Commented reference in settings.yaml under embedding.