diff --git a/test/collection/test_classes_generative.py b/test/collection/test_classes_generative.py index 5d9496417..db8b95144 100644 --- a/test/collection/test_classes_generative.py +++ b/test/collection/test_classes_generative.py @@ -106,6 +106,7 @@ def test_generative_parameters_images_parsing( target_model="arn:aws:sagemaker:us-west-2:123456789012:model/text-to-image", target_variant="variant-1", temperature=0.5, + stop_sequences=["\n"], )._to_grpc( _GenerativeConfigRuntimeOptions( return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"] @@ -121,6 +122,7 @@ def test_generative_parameters_images_parsing( target_model="arn:aws:sagemaker:us-west-2:123456789012:model/text-to-image", target_variant="variant-1", temperature=0.5, + stop_sequences=base_pb2.TextArray(values=["\n"]), images=base_pb2.TextArray(values=[LOGO_ENCODED]), image_properties=base_pb2.TextArray(values=["image"]), ), diff --git a/weaviate/collections/classes/generative.py b/weaviate/collections/classes/generative.py index d5c1a5de2..b6c5bed1f 100644 --- a/weaviate/collections/classes/generative.py +++ b/weaviate/collections/classes/generative.py @@ -130,9 +130,10 @@ def _to_grpc(self, opts: _GenerativeConfigRuntimeOptions) -> generative_pb2.Gene target_variant=self.target_variant, temperature=self.temperature, max_tokens=self.max_tokens, + stop_sequences=_to_text_array(self.stop_sequences), images=_to_text_array(opts.images), image_properties=_to_text_array(opts.image_properties), - # TODO - add top_k, top_p & stop_sequences here when added to server-side proto + # TODO - add top_k & top_p here when added to server-side proto # Check the latest availble version of `grpc/proto/v1/generative.proto` (see GenerativeAWS) in the server repo ), ) @@ -587,6 +588,7 @@ def aws( target_model: Optional[str] = None, target_variant: Optional[str] = None, temperature: Optional[float] = None, + stop_sequences: Optional[List[str]] = None, ) -> _GenerativeConfigRuntime: """Create a `_GenerativeAWS` object for use when performing dynamic AI generation using the `generative-aws` module. @@ -602,6 +604,7 @@ def aws( target_model: The target model to use. Defaults to `None`, which uses the server-defined default target_variant: The target variant to use. Defaults to `None`, which uses the server-defined default temperature: The temperature to use. Defaults to `None`, which uses the server-defined default + stop_sequences: The stop sequences to use. Defaults to `None`, which uses the server-defined default """ return _GenerativeAWS( model=model, @@ -616,7 +619,7 @@ def aws( temperature=temperature, top_k=None, top_p=None, - stop_sequences=None, + stop_sequences=stop_sequences, ) @staticmethod