Added Nova models on bedrock via converse.

Author: adivekar-utexas

src/fmcore/algorithm/bedrock.py

@@ -350,6 +350,44 @@ def call_llama_3(
         response_body: Dict = json.loads(response.get("body").read())
         return response_body.get("generation")
 
+    def call_nova(
+        bedrock_client,
+        model_name: str,
+        prompt: str,
+        max_tokens_to_sample: int,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        stop_sequences: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        messages = [
+            {
+                "role": "user",
+                "content": [{"text": prompt}],
+            }
+        ]
+
+        inference_config = {
+            "maxTokens": max_tokens_to_sample,
+        }
+        if temperature is not None:
+            assert isinstance(temperature, (float, int)) and 0 <= temperature <= 1
+            inference_config["temperature"] = temperature
+        if top_p is not None:
+            assert isinstance(top_p, (float, int)) and 0 <= top_p <= 1
+            inference_config["topP"] = top_p
+        if stop_sequences is not None:
+            assert isinstance(stop_sequences, list)
+            if len(stop_sequences) > 0:
+                inference_config["stopSequences"] = stop_sequences
+
+        response_body: Dict = bedrock_client.converse(
+            modelId=model_name,
+            messages=messages,
+            inferenceConfig=inference_config,
+        )
+        return "\n".join([d["text"] for d in response_body["output"]["message"]["content"]])
+
     def call_mistral(
         bedrock_client,
         model_name: str,
@@ -458,6 +496,14 @@ def call_bedrock(
                 **generation_params,
             )
             return generated_text
+        elif "nova" in model_name:
+            generated_text: str = call_nova(
+                bedrock_client=bedrock_client,
+                prompt=prompt,
+                model_name=model_name,
+                **generation_params,
+            )
+            return generated_text
         else:
             bedrock_invoke_model_params = {"prompt": prompt, **generation_params}
             response = bedrock_client.invoke_model(