Add logprobs to SamplingParameters in vllm (#9010)

This PR aims at amending #8806 , that I opened a few days ago, adding
the extra `logprobs` parameter that I accidentally forgot
This commit is contained in:
Massimiliano Pronesti 2023-08-10 00:48:29 +02:00 committed by GitHub
parent 8dd071ad08
commit c72da53c10
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -54,6 +54,9 @@ class VLLM(BaseLLM):
max_new_tokens: int = 512 max_new_tokens: int = 512
"""Maximum number of tokens to generate per output sequence.""" """Maximum number of tokens to generate per output sequence."""
logprobs: Optional[int] = None
"""Number of log probabilities to return per output token."""
client: Any #: :meta private: client: Any #: :meta private:
@root_validator() @root_validator()
@ -91,6 +94,7 @@ class VLLM(BaseLLM):
"stop": self.stop, "stop": self.stop,
"ignore_eos": self.ignore_eos, "ignore_eos": self.ignore_eos,
"use_beam_search": self.use_beam_search, "use_beam_search": self.use_beam_search,
"logprobs": self.logprobs,
} }
def _generate( def _generate(