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

improve(vertex_ai.py): Add frequency_penalty and presence_penalty.

Browse source
This commit is contained in:
David Manouchehri 2024-04-22 17:46:43 +00:00
parent 1a7eec5786
commit c643e04ada
No known key found for this signature in database
1 changed files with 16 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

16
litellm/llms/vertex_ai.py
View file

@ -34,6 +34,8 @@ class ExtendedGenerationConfig(dict):
max_output_tokens: Optional[int] = None, max_output_tokens: Optional[int] = None,
stop_sequences: Optional[List[str]] = None, stop_sequences: Optional[List[str]] = None,
response_mime_type: Optional[str] = None, response_mime_type: Optional[str] = None,
frequency_penalty: Optional[float] = None,
presence_penalty: Optional[float] = None,
): ):
super().__init__( super().__init__(
temperature=temperature, temperature=temperature,
@ -43,6 +45,8 @@ class ExtendedGenerationConfig(dict):
max_output_tokens=max_output_tokens, max_output_tokens=max_output_tokens,
stop_sequences=stop_sequences, stop_sequences=stop_sequences,
response_mime_type=response_mime_type, response_mime_type=response_mime_type,
frequency_penalty=frequency_penalty,
presence_penalty=presence_penalty,
) )
class VertexAIConfig: class VertexAIConfig:
@ -66,6 +70,10 @@ class VertexAIConfig:
- `stop_sequences` (List[str]): The set of character sequences (up to 5) that will stop output generation. If specified, the API will stop at the first appearance of a stop sequence. The stop sequence will not be included as part of the response. - `stop_sequences` (List[str]): The set of character sequences (up to 5) that will stop output generation. If specified, the API will stop at the first appearance of a stop sequence. The stop sequence will not be included as part of the response.
- `frequency_penalty` (float): This parameter is used to penalize the model from repeating the same output. The default value is 0.0.
- `presence_penalty` (float): This parameter is used to penalize the model from generating the same output as the input. The default value is 0.0.
Note: Please make sure to modify the default parameters as required for your use case. Note: Please make sure to modify the default parameters as required for your use case.
""" """
@ -76,6 +84,8 @@ class VertexAIConfig:
response_mime_type: Optional[str] = None response_mime_type: Optional[str] = None
candidate_count: Optional[int] = None candidate_count: Optional[int] = None
stop_sequences: Optional[list] = None stop_sequences: Optional[list] = None
frequency_penalty: Optional[float] = None
presence_penalty: Optional[float] = None
def __init__( def __init__(
self, self,
@ -86,6 +96,8 @@ class VertexAIConfig:
response_mime_type: Optional[str] = None, response_mime_type: Optional[str] = None,
candidate_count: Optional[int] = None, candidate_count: Optional[int] = None,
stop_sequences: Optional[list] = None, stop_sequences: Optional[list] = None,
frequency_penalty: Optional[float] = None,
presence_penalty: Optional[float] = None,
) -> None: ) -> None:
locals_ = locals() locals_ = locals()
for key, value in locals_.items(): for key, value in locals_.items():
@ -142,6 +154,10 @@ class VertexAIConfig:
optional_params["max_output_tokens"] = value optional_params["max_output_tokens"] = value
if param == "response_format" and value["type"] == "json_object": if param == "response_format" and value["type"] == "json_object":
optional_params["response_mime_type"] = "application/json" optional_params["response_mime_type"] = "application/json"
if param == "frequency_penalty":
optional_params["frequency_penalty"] = value
if param == "presence_penalty":
optional_params["presence_penalty"] = value
if param == "tools" and isinstance(value, list): if param == "tools" and isinstance(value, list):
from vertexai.preview import generative_models from vertexai.preview import generative_models