fix: missing request parameters

Author: meken

tests/unit/aiplatform/test_language_models.py

@@ -1775,6 +1775,8 @@ def test_text_generation_ga(self):
                 presence_penalty=1.0,
                 frequency_penalty=1.0,
                 logit_bias={1: 100.0, 2: -100.0},
+                seed=42,
+                echo=True,
             )
 
         expected_errors = (100,)
@@ -1788,6 +1790,8 @@ def test_text_generation_ga(self):
         assert prediction_parameters["presencePenalty"] == 1.0
         assert prediction_parameters["frequencyPenalty"] == 1.0
         assert prediction_parameters["logitBias"] == {1: 100.0, 2: -100.0}
+        assert prediction_parameters["seed"] == 42
+        assert prediction_parameters["echo"] is True
         assert response.text == _TEST_TEXT_GENERATION_PREDICTION["content"]
         assert response.errors == expected_errors
 

vertexai/language_models/_language_models.py

@@ -1289,6 +1289,8 @@ def predict(
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         logit_bias: Optional[Dict[int, float]] = None,
+        seed: Optional[int] = None,
+        echo: Optional[bool] = None,
     ) -> "MultiCandidateTextGenerationResponse":
         """Gets model response for a single prompt.
 
@@ -1321,6 +1323,14 @@ def predict(
                 Larger positive bias increases the probability of choosing the token.
                 Smaller negative bias decreases the probability of choosing the token.
                 Range: [-100.0, 100.0]
+            seed:
+                Decoder generates random noise with a pseudo random number generator, temperature * noise is added to
+                logits before sampling. The pseudo random number generator (prng) takes a seed as input, it generates
+                the same output with the same seed. If seed is not set, the seed used in decoder will not be
+                deterministic, thus the generated random noise will not be deterministic. If seed is set, the
+                generated random noise will be deterministic.
+            echo:
+                If true, the prompt is echoed in the generated text.
 
         Returns:
             A `MultiCandidateTextGenerationResponse` object that contains the text produced by the model.
@@ -1338,6 +1348,8 @@ def predict(
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             logit_bias=logit_bias,
+            seed=seed,
+            echo=echo,
         )
 
         prediction_response = self._endpoint.predict(
@@ -1370,6 +1382,8 @@ async def predict_async(
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         logit_bias: Optional[Dict[int, float]] = None,
+        seed: Optional[int] = None,
+        echo: Optional[bool] = None,
     ) -> "MultiCandidateTextGenerationResponse":
         """Asynchronously gets model response for a single prompt.
 
@@ -1402,6 +1416,14 @@ async def predict_async(
                 Larger positive bias increases the probability of choosing the token.
                 Smaller negative bias decreases the probability of choosing the token.
                 Range: [-100.0, 100.0]
+            seed:
+                Decoder generates random noise with a pseudo random number generator, temperature * noise is added to
+                logits before sampling. The pseudo random number generator (prng) takes a seed as input, it generates
+                the same output with the same seed. If seed is not set, the seed used in decoder will not be
+                deterministic, thus the generated random noise will not be deterministic. If seed is set, the
+                generated random noise will be deterministic.
+            echo:
+                If true, the prompt is echoed in the generated text.
 
         Returns:
             A `MultiCandidateTextGenerationResponse` object that contains the text produced by the model.
@@ -1419,6 +1441,8 @@ async def predict_async(
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             logit_bias=logit_bias,
+            seed=seed,
+            echo=echo,
         )
 
         prediction_response = await self._endpoint.predict_async(
@@ -1443,6 +1467,8 @@ def predict_streaming(
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         logit_bias: Optional[Dict[int, float]] = None,
+        seed: Optional[int] = None,
+        echo: Optional[bool] = None,
     ) -> Iterator[TextGenerationResponse]:
         """Gets a streaming model response for a single prompt.
 
@@ -1475,6 +1501,14 @@ def predict_streaming(
                 Larger positive bias increases the probability of choosing the token.
                 Smaller negative bias decreases the probability of choosing the token.
                 Range: [-100.0, 100.0]
+            seed:
+                Decoder generates random noise with a pseudo random number generator, temperature * noise is added to
+                logits before sampling. The pseudo random number generator (prng) takes a seed as input, it generates
+                the same output with the same seed. If seed is not set, the seed used in decoder will not be
+                deterministic, thus the generated random noise will not be deterministic. If seed is set, the
+                generated random noise will be deterministic.
+            echo:
+                If true, the prompt is echoed in the generated text.
 
         Yields:
             A stream of `TextGenerationResponse` objects that contain partial
@@ -1491,6 +1525,8 @@ def predict_streaming(
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             logit_bias=logit_bias,
+            seed=seed,
+            echo=echo,
         )
 
         prediction_service_client = self._endpoint._prediction_client
@@ -1521,6 +1557,8 @@ async def predict_streaming_async(
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         logit_bias: Optional[Dict[int, float]] = None,
+        seed: Optional[int] = None,
+        echo: Optional[bool] = None,
     ) -> AsyncIterator[TextGenerationResponse]:
         """Asynchronously gets a streaming model response for a single prompt.
 
@@ -1553,6 +1591,14 @@ async def predict_streaming_async(
                 Larger positive bias increases the probability of choosing the token.
                 Smaller negative bias decreases the probability of choosing the token.
                 Range: [-100.0, 100.0]
+            seed:
+                Decoder generates random noise with a pseudo random number generator, temperature * noise is added to
+                logits before sampling. The pseudo random number generator (prng) takes a seed as input, it generates
+                the same output with the same seed. If seed is not set, the seed used in decoder will not be
+                deterministic, thus the generated random noise will not be deterministic. If seed is set, the
+                generated random noise will be deterministic.
+            echo:
+                If true, the prompt is echoed in the generated text.
 
         Yields:
             A stream of `TextGenerationResponse` objects that contain partial
@@ -1569,6 +1615,8 @@ async def predict_streaming_async(
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             logit_bias=logit_bias,
+            seed=seed,
+            echo=echo,
         )
 
         prediction_service_async_client = self._endpoint._prediction_async_client
@@ -1605,6 +1653,8 @@ def _create_text_generation_prediction_request(
     presence_penalty: Optional[float] = None,
     frequency_penalty: Optional[float] = None,
     logit_bias: Optional[Dict[int, int]] = None,
+    seed: Optional[int] = None,
+    echo: Optional[bool] = None,
 ) -> "_PredictionRequest":
     """Prepares the text generation request for a single prompt.
 
@@ -1637,6 +1687,14 @@ def _create_text_generation_prediction_request(
             Larger positive bias increases the probability of choosing the token.
             Smaller negative bias decreases the probability of choosing the token.
             Range: [-100.0, 100.0]
+        seed:
+            Decoder generates random noise with a pseudo random number generator, temperature * noise is added to
+            logits before sampling. The pseudo random number generator (prng) takes a seed as input, it generates
+            the same output with the same seed. If seed is not set, the seed used in decoder will not be
+            deterministic, thus the generated random noise will not be deterministic. If seed is set, the
+            generated random noise will be deterministic.
+        echo:
+            If true, the prompt is echoed in the generated text.
 
     Returns:
         A `_PredictionRequest` object that contains prediction instance and parameters.
@@ -1683,6 +1741,12 @@ def _create_text_generation_prediction_request(
     if logit_bias is not None:
         prediction_parameters["logitBias"] = logit_bias
 
+    if seed is not None:
+        prediction_parameters["seed"] = seed
+
+    if echo is not None:
+        prediction_parameters["echo"] = echo
+
     return _PredictionRequest(
         instance=instance,
         parameters=prediction_parameters,