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
"""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:
@root_validator()
@ -91,6 +94,7 @@ class VLLM(BaseLLM):
"stop": self.stop,
"ignore_eos": self.ignore_eos,
"use_beam_search": self.use_beam_search,
"logprobs": self.logprobs,
}
def _generate(