Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import re
 
+import colorcet
 import numpy as np
 import panel as pn
 import tastymap as tm
@@ -8,11 +9,10 @@ from openai import AsyncOpenAI
 
 pn.extension()
 
-COLORMAP = "viridis_r"
-NUM_COLORS = 8
-
 
 def color_by_logprob(text, log_prob):
+    tmap = tm.cook_tmap(colormap_input.value, colors_input.value)
+    colors = tmap.to_model("hex")
     linear_prob = np.round(np.exp(log_prob) * 100, 2)
     # select index based on probability
     color_index = int(linear_prob // (100 / (len(colors) - 1)))
@@ -37,6 +37,7 @@ def custom_serializer(content):
 
 
 async def respond_to_input(contents: str, user: str, instance: pn.chat.ChatInterface):
+    print(contents)
     if api_key_input.value:
         aclient.api_key = api_key_input.value
     elif not os.environ["OPENAI_API_KEY"]:
@@ -91,13 +92,10 @@ async def respond_to_input(contents: str, user: str, instance: pn.chat.ChatInter
     )
 
 
-tmap = tm.cook_tmap(COLORMAP, NUM_COLORS)
-colors = tmap.to_model("hex")
 aclient = AsyncOpenAI()
 api_key_input = pn.widgets.PasswordInput(
     name="API Key",
     placeholder="sk-...",
-    width=150,
 )
 system_input = pn.widgets.TextAreaInput(
     name="System Prompt",
@@ -112,16 +110,15 @@ default_input = pn.widgets.TextAreaInput(
 model_selector = pn.widgets.Select(
     name="Model",
     options=["gpt-3.5-turbo", "gpt-4"],
-    width=150,
 )
+colormap_input = pn.widgets.ColorMap(name="ColorMap", options=colorcet.palette)
+colors_input = pn.widgets.IntInput(name="Number of Colors", start=1, value=6)
 temperature_input = pn.widgets.FloatInput(
-    name="Temperature", start=0, end=2, step=0.01, value=1, width=100
-)
-max_tokens_input = pn.widgets.IntInput(name="Max Tokens", start=0, value=256, width=100)
-seed_input = pn.widgets.IntInput(name="Seed", start=0, end=100, value=0, width=100)
-memory_toggle = pn.widgets.Toggle(
-    name="Include Memory", value=False, width=100, margin=(22, 5)
+    name="Temperature", start=0, end=2, step=0.01, value=1
 )
+max_tokens_input = pn.widgets.IntInput(name="Max Tokens", start=0, value=256)
+seed_input = pn.widgets.IntInput(name="Seed", start=0, end=100, value=0)
+memory_toggle = pn.widgets.Toggle(name="Include Memory", value=False, margin=(22, 5))
 chat_interface = pn.chat.ChatInterface(
     callback=respond_to_input,
     callback_user="ChatGPT",
@@ -131,24 +128,22 @@ chat_interface = pn.chat.ChatInterface(
     },
     show_rerun=False,
 )
-pn.Column(
-    pn.Column(
-        pn.Row(
-            api_key_input,
-            system_input,
-            default_input,
-            model_selector,
-            align="center",
-        ),
-        pn.Row(
-            temperature_input,
-            max_tokens_input,
-            seed_input,
-            memory_toggle,
-            align="center",
-        ),
-        align="center",
-    ),
-    pn.Row(tmap._repr_html_(), align="center"),
-    chat_interface,
-).servable()
+
+sidebar = pn.Column(
+    api_key_input,
+    system_input,
+    default_input,
+    model_selector,
+    colormap_input,
+    colors_input,
+    temperature_input,
+    max_tokens_input,
+    seed_input,
+    memory_toggle,
+)
+main = pn.Column(chat_interface)
+template = pn.template.FastListTemplate(
+    sidebar=sidebar,
+    main=main,
+)
+template.show()