mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
add logit bias
This commit is contained in:
parent
dc7d2b7631
commit
abfc286c20
2 changed files with 3 additions and 2 deletions
|
@ -77,7 +77,7 @@ def test_completion_claude_stream():
|
|||
def test_completion_cohere():
|
||||
try:
|
||||
response = completion(
|
||||
model="command-nightly", messages=messages, max_tokens=100
|
||||
model="command-nightly", messages=messages, max_tokens=100, logit_bias={40: 10}
|
||||
)
|
||||
# Add any assertions here to check the response
|
||||
print(response)
|
||||
|
@ -91,7 +91,6 @@ def test_completion_cohere():
|
|||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
def test_completion_cohere_stream():
|
||||
try:
|
||||
messages = [
|
||||
|
|
|
@ -452,6 +452,8 @@ def get_optional_params(
|
|||
optional_params["temperature"] = temperature
|
||||
if max_tokens != float("inf"):
|
||||
optional_params["max_tokens"] = max_tokens
|
||||
if logit_bias != {}:
|
||||
optional_params["logit_bias"] = logit_bias
|
||||
return optional_params
|
||||
elif custom_llm_provider == "replicate":
|
||||
# any replicate models
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue