phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add logit bias

Browse source
This commit is contained in:
ishaan-jaff 2023-08-19 16:29:12 -07:00
parent dc7d2b7631
commit abfc286c20
2 changed files with 3 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/tests/test_completion.py
View file

@ -77,7 +77,7 @@ def test_completion_claude_stream():
def test_completion_cohere(): def test_completion_cohere():
try: try:
response = completion( 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 # Add any assertions here to check the response
print(response) print(response)
@ -91,7 +91,6 @@ def test_completion_cohere():
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
def test_completion_cohere_stream(): def test_completion_cohere_stream():
try: try:
messages = [ messages = [

2
litellm/utils.py
View file

@ -452,6 +452,8 @@ def get_optional_params(
optional_params["temperature"] = temperature optional_params["temperature"] = temperature
if max_tokens != float("inf"): if max_tokens != float("inf"):
optional_params["max_tokens"] = max_tokens optional_params["max_tokens"] = max_tokens
if logit_bias != {}:
optional_params["logit_bias"] = logit_bias
return optional_params return optional_params
elif custom_llm_provider == "replicate": elif custom_llm_provider == "replicate":
# any replicate models # any replicate models