fix: Correct logit_bias type annotation to accept keys as strings

sasha-gitg · copybara-github · commit 2676d25e62d1 · 2024-07-23T14:35:08.000-07:00
PiperOrigin-RevId: 655305894
diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
@@ -1354,7 +1354,7 @@ def predict(
         logprobs: Optional[int] = None,
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
-        logit_bias: Optional[Dict[int, float]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
         seed: Optional[int] = None,
     ) -> "MultiCandidateTextGenerationResponse":
         """Gets model response for a single prompt.
@@ -1443,7 +1443,7 @@ async def predict_async(
         logprobs: Optional[int] = None,
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
-        logit_bias: Optional[Dict[int, float]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
         seed: Optional[int] = None,
     ) -> "MultiCandidateTextGenerationResponse":
         """Asynchronously gets model response for a single prompt.
@@ -1524,7 +1524,7 @@ def predict_streaming(
         logprobs: Optional[int] = None,
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
-        logit_bias: Optional[Dict[int, float]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
         seed: Optional[int] = None,
     ) -> Iterator[TextGenerationResponse]:
         """Gets a streaming model response for a single prompt.
@@ -1610,7 +1610,7 @@ async def predict_streaming_async(
         logprobs: Optional[int] = None,
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
-        logit_bias: Optional[Dict[int, float]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
         seed: Optional[int] = None,
     ) -> AsyncIterator[TextGenerationResponse]:
         """Asynchronously gets a streaming model response for a single prompt.
@@ -1702,7 +1702,7 @@ def _create_text_generation_prediction_request(
     logprobs: Optional[int] = None,
     presence_penalty: Optional[float] = None,
     frequency_penalty: Optional[float] = None,
-    logit_bias: Optional[Dict[int, int]] = None,
+    logit_bias: Optional[Dict[str, float]] = None,
     seed: Optional[int] = None,
 ) -> "_PredictionRequest":
     """Prepares the text generation request for a single prompt.
