Update app.py

app.py

@@ -1,10 +1,11 @@
 import gradio as gr
-#from huggingface_hub import InferenceClient
 import random
 import spaces
 import os
 
 os.environ["KERAS_BACKEND"] = "jax"
+os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"
+
 import keras_hub
 import keras
 
@@ -29,9 +30,8 @@ def format_prompt(message, history):
     return prompt
 
 
-def chat_inf(system_prompt, prompt, history, tokens, client_choice, temp, top_p, seed):
+def chat_inf(system_prompt, prompt, history, client_choice, seed, temp, tokens, top_k, top_p):
     client = clients[int(client_choice) - 1]
-    
     if not history:
         history = []
         hist_len = 0
@@ -39,7 +39,7 @@ def chat_inf(system_prompt, prompt, history, tokens, client_choice, temp, top_p,
         hist_len = len(history)
         print(hist_len)
 
-    sampler = keras_hub.samplers.TopPSampler(p=top_p, seed=seed, temperature=temp)
+    sampler = keras_nlp.samplers.TopKSampler(k=top_k, seed=seed, temperature=temp)
     client.compile(sampler=sampler)
     
     formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
@@ -47,7 +47,7 @@ def chat_inf(system_prompt, prompt, history, tokens, client_choice, temp, top_p,
     output = ""
 
     for response in stream:
-        output += response.str(tokens)
+        output += response.token.text
         yield [(prompt, output)]
     history.append((prompt, output))
     yield history
@@ -67,14 +67,14 @@ def check_rand(inp, val):
         return gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=int(val))
 
 
-with gr.Blocks(theme=gr.themes.Glass(),css=".gradio-container {background-color: rgb(134 239 172)}") as app:
+with gr.Blocks(css=".gradio-container {background-color: rgb(187 247 208)}", gr.themes.Monochrome()) as app:
     gr.HTML(
-        """<center><h1 style='font-size:xx-large;'>CalmChat</h1></center>""")
+        """<center><h1 style='font-size:xx-large;'>CalmChat: A Mental Health Conversational Agent</h1></center>""")
     with gr.Group():
         with gr.Row():
             client_choice = gr.Dropdown(label="Models", type='index', choices=[c for c in models], value=models[0],
                                         interactive=True)
-    chat_b = gr.Chatbot(height=500)
+    chat_b = gr.Chatbot(height=300)
     with gr.Group():
         with gr.Row():
             with gr.Column(scale=1):
@@ -86,8 +86,8 @@ with gr.Blocks(theme=gr.themes.Glass(),css=".gradio-container {background-color:
             with gr.Column(scale=1):
                 with gr.Group():
                     temp = gr.Slider(label="Temperature", step=0.01, minimum=0.01, maximum=1.0, value=0.9)
+                    top_k = gr.Slider(label="Top-K", step=0.5, minimum=1, maximum=10, value=5)
                     top_p = gr.Slider(label="Top-P", step=0.01, minimum=0.01, maximum=1.0, value=0.9)
-                
 
     with gr.Group():
         with gr.Row():
@@ -99,9 +99,12 @@ with gr.Blocks(theme=gr.themes.Glass(),css=".gradio-container {background-color:
                     stop_btn = gr.Button("Stop")
                     clear_btn = gr.Button("Clear")
 
-    chat_sub = inp.submit(check_rand, [rand, seed],seed).then(chat_inf,[sys_inp, inp, chat_b, client_choice, tokens], chat_b, temp, top_p, seed)
-    go = btn.click(check_rand, [rand, seed], seed).then(chat_inf,[sys_inp, inp, chat_b, client_choice, tokens], chat_b, temp, top_p, seed)
-    
+    chat_sub = inp.submit(check_rand, [rand, seed], seed).then(chat_inf,
+                                                               [sys_inp, inp, chat_b, client_choice, seed, temp, tokens,
+                                                                top_p, top_k], chat_b)
+    go = btn.click(check_rand, [rand, seed], seed).then(chat_inf,
+                                                        [sys_inp, inp, chat_b, client_choice, seed, temp, tokens, top_p,
+                                                         top_k], chat_b)
     stop_btn.click(None, None, None, cancels=[go, chat_sub])
     clear_btn.click(clear_fn, None, [chat_b])
 app.queue(default_concurrency_limit=10).launch()