feat: Bedrock model id support for jamba-1.5 (#212)

* feat: Bedrock model id support for jamba-1.5

* fix: Bedrock examples

Author: Josephasafg

README.md

@@ -560,7 +560,7 @@ messages = [
 
 response = client.chat.completions.create(
     messages=messages,
-    model_id=BedrockModelID.JAMBA_INSTRUCT_V1,
+    model_id=BedrockModelID.JAMBA_1_5_LARGE,
 )
 ```
 
@@ -582,7 +582,7 @@ client = AI21BedrockClient()
 
 response = client.chat.completions.create(
     messages=messages,
-    model=BedrockModelID.JAMBA_INSTRUCT_V1,
+    model=BedrockModelID.JAMBA_1_5_LARGE,
     stream=True,
 )
 
@@ -607,7 +607,7 @@ messages = [
 async def main():
     response = await client.chat.completions.create(
         messages=messages,
-        model_id=BedrockModelID.JAMBA_INSTRUCT_V1,
+        model_id=BedrockModelID.JAMBA_1_5_LARGE,
     )
 
 
@@ -633,7 +633,7 @@ messages = [
 
 response = client.chat.completions.create(
     messages=messages,
-    model_id=BedrockModelID.JAMBA_INSTRUCT_V1,
+    model_id=BedrockModelID.JAMBA_1_5_LARGE,
 )
 ```
 
@@ -658,7 +658,7 @@ messages = [
 async def main():
   response = await client.chat.completions.create(
       messages=messages,
-      model_id=BedrockModelID.JAMBA_INSTRUCT_V1,
+      model_id=BedrockModelID.JAMBA_1_5_LARGE,
   )
 
 asyncio.run(main())

ai21/clients/bedrock/bedrock_model_id.py

@@ -2,3 +2,5 @@ class BedrockModelID:
     J2_MID_V1 = "ai21.j2-mid-v1"
     J2_ULTRA_V1 = "ai21.j2-ultra-v1"
     JAMBA_INSTRUCT_V1 = "ai21.jamba-instruct-v1:0"
+    JAMBA_1_5_MINI = "ai21.jamba-1-5-mini-v1:0"
+    JAMBA_1_5_LARGE = "ai21.jamba-1-5-large-v1:0"

examples/bedrock/chat/async_chat_completions.py

@@ -13,7 +13,7 @@
 async def main():
     response = await client.chat.completions.create(
         messages=messages,
-        model=BedrockModelID.JAMBA_INSTRUCT_V1,
+        model=BedrockModelID.JAMBA_1_5_MINI,
     )
 
     print(f"response: {response}")

examples/bedrock/chat/async_stream_chat_completions.py

@@ -16,7 +16,7 @@
 async def main():
     response = await client.chat.completions.create(
         messages=messages,
-        model=BedrockModelID.JAMBA_INSTRUCT_V1,
+        model=BedrockModelID.JAMBA_1_5_MINI,
         max_tokens=100,
         stream=True,
     )

examples/bedrock/chat/chat_completions.py

@@ -22,7 +22,7 @@
     messages=messages,
     max_tokens=1000,
     temperature=0,
-    model=BedrockModelID.JAMBA_INSTRUCT_V1,
+    model=BedrockModelID.JAMBA_1_5_MINI,
 )
 
 print(f"response: {response}")

examples/bedrock/chat/stream_chat_completions.py

@@ -13,7 +13,7 @@
 
 response = client.chat.completions.create(
     messages=messages,
-    model=BedrockModelID.JAMBA_INSTRUCT_V1,
+    model=BedrockModelID.JAMBA_1_5_MINI,
     max_tokens=100,
     stream=True,
 )