From bb8c095127bffcb48a18744bfe2dfd53223814af Mon Sep 17 00:00:00 2001 From: Louis Date: Mon, 4 Sep 2023 18:53:48 +0100 Subject: [PATCH] Add 'download_dir' argument to VLLM (#9754) - Description: Add a 'download_dir' argument to VLLM model (to change the cache download directotu when retrieving a model from HF hub) - Issue: On some remote machine, I want the cache dir to be in a volume where I have space (models are heavy nowadays). Sometimes the default HF cache dir might not be what we want. - Dependencies: None --------- Co-authored-by: Harrison Chase --- libs/langchain/langchain/llms/vllm.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/libs/langchain/langchain/llms/vllm.py b/libs/langchain/langchain/llms/vllm.py index 1a6e1a5910b..537a9bbb6f3 100644 --- a/libs/langchain/langchain/llms/vllm.py +++ b/libs/langchain/langchain/llms/vllm.py @@ -62,6 +62,10 @@ class VLLM(BaseLLM): dtype: str = "auto" """The data type for the model weights and activations.""" + download_dir: Optional[str] = None + """Directory to download and load the weights. (Default to the default + cache dir of huggingface)""" + vllm_kwargs: Dict[str, Any] = Field(default_factory=dict) """Holds any model parameters valid for `vllm.LLM` call not explicitly specified.""" @@ -84,6 +88,7 @@ class VLLM(BaseLLM): tensor_parallel_size=values["tensor_parallel_size"], trust_remote_code=values["trust_remote_code"], dtype=values["dtype"], + download_dir=values["download_dir"], **values["vllm_kwargs"], )