add command-r plus

.gitignore

@@ -3,6 +3,7 @@ gradio_cached_examples/
 static/*
 static/
 notes.txt
+test_model.py
 
 *.py[cod]
 

app.py

@@ -257,6 +257,62 @@ def call_claude_3_haiku(system_prompt, prompt, temperature):
     return results
 
 
+def call_command_r_plus(system_prompt, prompt, temperature):
+
+    chat_history = [
+        {"role": "USER", "message": system_prompt},
+    ]
+
+    prompt_config = {
+        "message": prompt,
+        "max_tokens": 4096,
+        "chat_history": chat_history,
+        "temperature": temperature,
+    }
+
+    body = json.dumps(prompt_config)
+
+    modelId = "cohere.command-r-plus-v1:0"
+    accept = "application/json"
+    contentType = "application/json"
+
+    response = bedrock_runtime.invoke_model(
+        body=body, modelId=modelId, accept=accept, contentType=contentType
+    )
+    response_body = json.loads(response.get("body").read())
+
+    results = response_body.get("text")
+    return results
+
+
+def call_command_r(system_prompt, prompt, temperature):
+
+    chat_history = [
+        {"role": "USER", "message": system_prompt},
+    ]
+
+    prompt_config = {
+        "message": prompt,
+        "max_tokens": 4096,
+        "chat_history": chat_history,
+        "temperature": temperature,
+    }
+
+    body = json.dumps(prompt_config)
+
+    modelId = "cohere.command-r-v1:0"
+    accept = "application/json"
+    contentType = "application/json"
+
+    response = bedrock_runtime.invoke_model(
+        body=body, modelId=modelId, accept=accept, contentType=contentType
+    )
+    response_body = json.loads(response.get("body").read())
+
+    results = response_body.get("text")
+    return results
+
+
 system_prompt_text = """
 As an esteemed level designer renowned for creating some of the top 100 levels in Super Mario Maker, you are tasked with crafting a playable section for the original Super Mario on NES. Your extensive experience and creativity are key to designing levels that are not only challenging but also immensely enjoyable. Use the following symbols to represent different game elements, ensuring each level is a masterpiece of design:
 
@@ -403,6 +459,10 @@ def generate(model, prompt, temperature, system_prompt=system_prompt_text):
         level = call_llama3_8b(system_prompt, prompt, temperature)
     elif model == "Claude Opus":
         level = call_claude_3_opus(system_prompt, prompt, temperature)
+    elif model == "Command R Plus":
+        level = call_command_r_plus(system_prompt, prompt, temperature)
+    elif model == "Command R":
+        level = call_command_r(system_prompt, prompt, temperature)
     else:
         raise ValueError("Invalid model")
 
@@ -449,8 +509,8 @@ You can try it out by entering in a prompt and clicking `Generate level` to play
                 )
 
             text_prompt = gr.Textbox(
-                value="Generate a level with a few pipes, many coins. make sure there are only 10 enemies. Make sure there is a ground path Mario can walk on",
-                label="Enter your MarioGPT prompt. ex: 'Generate a level with a few pipes, many coins. make sure there are only 10 enemies. Make sure there is a ground path Mario can walk on'",
+                value="Generate a level with a few pipes, many coins. make sure there are around 10 enemies. Make sure there is a ground path Mario can walk on",
+                label="Enter your MarioGPT prompt. ex: 'Generate a level with a few pipes, many coins. make sure there are around 10 enemies. Make sure there is a ground path Mario can walk on'",
             )
 
             model = gr.Radio(
@@ -460,6 +520,8 @@ You can try it out by entering in a prompt and clicking `Generate level` to play
                     "Claude Haiku",
                     "Llama3 70B",
                     "Llama3 8B",
+                    "Command R Plus",
+                    # "Command R", # taking too long to generate
                 ],
                 label="Select Model",
                 value="Claude Opus",
@@ -510,7 +572,7 @@ You can try it out by entering in a prompt and clicking `Generate level` to play
                 0.7,
             ],
             [
-                "Claude Sonnet",
+                "Command R Plus",
                 "Make a simple level that has no enemies, but lots and lots of coins. Lots of blocks for Mario to walk on.",
                 0.7,
             ],