Add support for logit_bias and logit_bias_type parameters

This commit is contained in:
Tanner Hobson
2023-06-09 13:13:08 -04:00
parent 0da655b3be
commit eb7645b3ba
2 changed files with 53 additions and 2 deletions

View File

@@ -1380,6 +1380,7 @@ class Llama:
mirostat_tau: float = 5.0,
mirostat_eta: float = 0.1,
model: Optional[str] = None,
logits_processor: Optional[LogitsProcessorList] = None,
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
"""Generate a chat completion from a list of messages.
@@ -1421,6 +1422,7 @@ class Llama:
mirostat_tau=mirostat_tau,
mirostat_eta=mirostat_eta,
model=model,
logits_processor=logits_processor,
)
if stream:
chunks: Iterator[CompletionChunk] = completion_or_chunks # type: ignore