From b775d15f2e5f343c0dc6d5eaf97749ab4c2e93f1 Mon Sep 17 00:00:00 2001
From: Yuval <yuval.belfer@gmail.com>
Date: Sun, 2 Jul 2023 07:10:56 +0300
Subject: [PATCH] fix: Updated AI21 models from J1 to J2 (#104)

also, added ai21 example to README

Co-authored-by: Yuval Belfer <yuvalb@ai21.com>
---
 README.md                |  7 +++++++
 manifest/clients/ai21.py | 14 +++++++-------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 34e7715..175d38f 100644
--- a/README.md
+++ b/README.md
@@ -78,6 +78,13 @@ manifest = Manifest(
 ```
 You can also just set `export COHERE_API_KEY=<COHERE_API_KEY>` and not use `client_connection`.
 
+If you want to use AI21 Labs, run
+```python
+manifest = Manifest(
+    client_name = "ai21",
+    client_connection = <AI21_API_KEY>,
+)
+```
 
 You can see the model details and possible model inputs to `run()` via
 ```python
diff --git a/manifest/clients/ai21.py b/manifest/clients/ai21.py
index 03a8669..8db5b58 100644
--- a/manifest/clients/ai21.py
+++ b/manifest/clients/ai21.py
@@ -9,9 +9,9 @@ from manifest.request import LMRequest
 logger = logging.getLogger(__name__)
 
 AI21_ENGINES = {
-    "j1-jumbo",
-    "j1-grande",
-    "j1-large",
+    "j2-ultra",
+    "j2-mid",
+    "j2-light",
 }
 
 
@@ -20,9 +20,9 @@ class AI21Client(Client):
 
     # User param -> (client param, default value)
     PARAMS = {
-        "engine": ("engine", "j1-large"),
-        "temperature": ("temperature", 1.0),
-        "max_tokens": ("maxTokens", 10),
+        "engine": ("engine", "j2-ultra"),
+        "temperature": ("temperature", 0.7),
+        "max_tokens": ("maxTokens", 40),
         "top_k": ("topKReturn", 0),
         "n": ("numResults", 1),
         "top_p": ("topP", 1.0),
@@ -45,7 +45,7 @@ class AI21Client(Client):
             connection_str: connection string.
             client_args: client arguments.
         """
-        # Taken from https://studio.ai21.com/docs/api/
+        # Taken from https://docs.ai21.com/
         self.host = "https://api.ai21.com/studio/v1"
         self.api_key = connection_str or os.environ.get("AI21_API_KEY")
         if self.api_key is None: