Allow specifying arbitrary keyword arguments in langchain.llms.VLLM (#9683)

Description: add arbitrary keyword arguments for VLLM
Issue: https://github.com/langchain-ai/langchain/issues/9682
Dependencies: none
Tag maintainer: @hwchase17, @baskaryan
This commit is contained in:
Jayson Ng 2023-09-03 23:40:06 -04:00 committed by GitHub
parent 43c4c6dfcc
commit 68f2363f5d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -3,7 +3,7 @@ from typing import Any, Dict, List, Optional
from langchain.callbacks.manager import CallbackManagerForLLMRun from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import BaseLLM from langchain.llms.base import BaseLLM
from langchain.llms.openai import BaseOpenAI from langchain.llms.openai import BaseOpenAI
from langchain.pydantic_v1 import root_validator from langchain.pydantic_v1 import Field, root_validator
from langchain.schema.output import Generation, LLMResult from langchain.schema.output import Generation, LLMResult
@ -62,6 +62,9 @@ class VLLM(BaseLLM):
dtype: str = "auto" dtype: str = "auto"
"""The data type for the model weights and activations.""" """The data type for the model weights and activations."""
vllm_kwargs: Dict[str, Any] = Field(default_factory=dict)
"""Holds any model parameters valid for `vllm.LLM` call not explicitly specified."""
client: Any #: :meta private: client: Any #: :meta private:
@root_validator() @root_validator()
@ -81,6 +84,7 @@ class VLLM(BaseLLM):
tensor_parallel_size=values["tensor_parallel_size"], tensor_parallel_size=values["tensor_parallel_size"],
trust_remote_code=values["trust_remote_code"], trust_remote_code=values["trust_remote_code"],
dtype=values["dtype"], dtype=values["dtype"],
**values["vllm_kwargs"],
) )
return values return values