Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 from threading import Thread
 from typing import Iterator
-
 import gradio as gr
 import spaces
 import torch
@@ -9,8 +8,36 @@ import json
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 DESCRIPTION = """\
-Shakti is a 2.5 billion parameter language model specifically optimized for resource-constrained environments such as edge devices, including smartphones, wearables, and IoT systems. With support for vernacular languages and domain-specific tasks, Shakti excels in industries such as healthcare, finance, and customer service.
-For more details, please check [here](https://arxiv.org/pdf/2410.11331v1).
+Shakti LLMs (Large Language Models) are a group of compact language models specifically optimized for resource-constrained environments such as edge devices, including smartphones, wearables, and IoT (Internet of Things) systems. These models provide support for vernacular languages and domain-specific tasks, making them particularly suitable for industries such as healthcare, finance, and customer service.
+For more details, please check [here](https://arxiv.org/pdf/2410.11331v1)
+"""
+
+
+# """\
+# Shakti LLMs are a group of small language model specifically optimized for resource-constrained environments such as edge devices, including smartphones, wearables, and IoT systems. With support for vernacular languages and domain-specific tasks, Shakti excels in industries such as healthcare, finance, and customer service.
+# For more details, please check [here](https://arxiv.org/pdf/2410.11331v1).
+# """
+
+
+# Custom CSS for the send button
+CUSTOM_CSS = """
+.send-btn {
+    padding: 0.5rem !important;
+    width: 55px !important;
+    height: 55px !important;
+    border-radius: 50% !important;
+    margin-top: 1rem;
+    cursor: pointer;
+}
+
+.send-btn svg {
+    width: 20px !important;
+    height: 20px !important;
+    position: absolute;
+    top: 50%;
+    left: 50%;
+    transform: translate(-50%, -50%);
+}
 """
 
 MAX_MAX_NEW_TOKENS = 2048
@@ -43,14 +70,14 @@ def load_model(selected_model: str):
         token=os.getenv("SHAKTI")
     )
     model.eval()
-    current_model = selected_model  # Update the current model
+    print("Selected Model: ", selected_model)
+    current_model = selected_model
 
 
-# Initial model load (default to 2.5B)
+# Initial model load
 load_model("Shakti-2.5B")
 
 
-@spaces.GPU(duration=90)
 def generate(
         message: str,
         chat_history: list[tuple[str, str]],
@@ -62,24 +89,19 @@ def generate(
 ) -> Iterator[str]:
     conversation = []
 
-    # Conditional logic for adding prompt based on model
     if current_model == "Shakti-2.5B":
         for user, assistant in chat_history:
-            conversation.extend(
-                [
-                    json.loads(os.getenv("PROMPT")),
-                    {"role": "user", "content": user},
-                    {"role": "assistant", "content": assistant},
-                ]
-            )
+            conversation.extend([
+                json.loads(os.getenv("PROMPT")),
+                {"role": "user", "content": user},
+                {"role": "assistant", "content": assistant},
+            ])
     else:
         for user, assistant in chat_history:
-            conversation.extend(
-                [
-                    {"role": "user", "content": user},
-                    {"role": "assistant", "content": assistant},
-                ]
-            )
+            conversation.extend([
+                {"role": "user", "content": user},
+                {"role": "assistant", "content": assistant},
+            ])
 
     conversation.append({"role": "user", "content": message})
 
@@ -110,72 +132,167 @@ def generate(
         yield "".join(outputs)
 
 
-def update_examples(selected_model):
-    if selected_model == "Shakti-100M":
-        return [["Tell me a story"],
-                ["Write a short poem on Rose"],
-                ["What are computers"]]
-    elif selected_model == "Shakti-250M":
-        return [["Can you explain the pathophysiology of hypertension and its impact on the cardiovascular system?"],
-                ["What are the potential side effects of beta-blockers in the treatment of arrhythmias?"],
-                ["What foods are good for boosting the immune system?"],
-                ["What is the difference between a stock and a bond?"],
-                ["How can I start saving for retirement?"],
-                ["What are some low-risk investment options?"],
-                ["What is a power of attorney and when is it used?"],
-                ["What are the key differences between a will and a trust?"],
-                ["How do I legally protect my business name?"]]
-    else:
-        return [["Tell me a story"], ["write a short poem which is hard to sing"],
-                ['मुझे भारतीय इतिहास के बारे में बताएं']]
+def respond(message, chat_history, max_new_tokens, temperature):
+    bot_message = ""
+    for chunk in generate(message, chat_history, max_new_tokens, temperature):
+        bot_message += chunk
+    chat_history.append((message, bot_message))
+    return "", chat_history
+
+
+def get_examples(selected_model):
+    examples = {
+        "Shakti-100M": [
+            ["Tell me a story"],
+            ["Write a short poem on Rose"],
+            ["What are computers"]
+        ],
+        "Shakti-250M": [
+            ["Can you explain the pathophysiology of hypertension and its impact on the cardiovascular system?"],
+            ["What are the potential side effects of beta-blockers in the treatment of arrhythmias?"],
+            ["What foods are good for boosting the immune system?"],
+            ["What is the difference between a stock and a bond?"],
+            ["How can I start saving for retirement?"],
+            ["What are some low-risk investment options?"]
+        ],
+        "Shakti-2.5B": [
+            ["Tell me a story"],
+            ["write a short poem which is hard to sing"],
+            ['मुझे भारतीय इतिहास के बारे में बताएं']
+        ]
+    }
+    return examples.get(selected_model, [])
 
 
 def on_model_select(selected_model):
     load_model(selected_model)  # Load the selected model
-    examples = update_examples(selected_model)  # Update examples
-    return gr.update(examples=examples), gr.update(value=[])  # Clear the chat space and update examples
+    # Return the message and chat history updates
+    return gr.update(value=""), gr.update(value=[])  # Clear message and chat history
+
+
+def update_examples_visibility(selected_model):
+    # Return individual updates for each example section
+    return (
+        gr.update(visible=selected_model == "Shakti-100M"),
+        gr.update(visible=selected_model == "Shakti-250M"),
+        gr.update(visible=selected_model == "Shakti-2.5B")
+    )
 
 
-chat_history = gr.Chatbot()
+def example_selector(example):
+    return example
 
-with gr.Blocks(css="style.css", fill_height=True) as demo:
+
+with gr.Blocks(css=CUSTOM_CSS) as demo:
     gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
-
-    # Dropdown for model selection
-    model_dropdown = gr.Dropdown(
-        label="Select Model",
-        choices=["Shakti-100M", "Shakti-250M", "Shakti-2.5B"],
-        value="Shakti-2.5B",
-        interactive=True,
-    )
 
-    # Create the interface with dynamic inputs and chat history
-    max_tokens_slider = gr.Slider(
-        label="Max new tokens",
-        minimum=1,
-        maximum=MAX_MAX_NEW_TOKENS,
-        step=1,
-        value=DEFAULT_MAX_NEW_TOKENS,
-    )
+    with gr.Row():
+        model_dropdown = gr.Dropdown(
+            label="Select Model",
+            choices=list(model_options.keys()),
+            value="Shakti-2.5B",
+            interactive=True
+        )
 
-    temperature_slider = gr.Slider(
-        label="Temperature",
-        minimum=0.1,
-        maximum=4.0,
-        step=0.1,
-        value=0.6,
+    chatbot = gr.Chatbot()
+
+    with gr.Row():
+        with gr.Column(scale=20):
+            msg = gr.Textbox(
+                label="Message",
+                placeholder="Enter your message here",
+                lines=2,
+                show_label=False
+            )
+        with gr.Column(scale=1, min_width=50):
+            send_btn = gr.Button(
+                value="➤",
+                variant="primary",
+                elem_classes=["send-btn"]
+            )
+
+    with gr.Accordion("Parameters", open=False):
+        max_tokens_slider = gr.Slider(
+            label="Max new tokens",
+            minimum=1,
+            maximum=MAX_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        )
+        temperature_slider = gr.Slider(
+            label="Temperature",
+            minimum=0.1,
+            maximum=4.0,
+            step=0.1,
+            value=0.6,
+        )
+
+    # Add submit action handlers
+    submit_click = send_btn.click(
+        respond,
+        inputs=[msg, chatbot, max_tokens_slider, temperature_slider],
+        outputs=[msg, chatbot]
     )
 
-    chat_interface = gr.Interface(
-        fn=generate,
-        inputs=[gr.Textbox(lines=2, placeholder="Enter your message here"), chat_history, max_tokens_slider,
-                temperature_slider],
-        outputs=chat_history,
-        live=True,
+    submit_enter = msg.submit(
+        respond,
+        inputs=[msg, chatbot, max_tokens_slider, temperature_slider],
+        outputs=[msg, chatbot]
     )
 
-    # Function to handle model change and update examples dynamically
-    model_dropdown.change(on_model_select, inputs=model_dropdown, outputs=[chat_interface, chat_history])
+    # Create separate example sections for each model
+    with gr.Row():
+        with gr.Column(visible=False) as examples_100m:
+            gr.Examples(
+                examples=get_examples("Shakti-100M"),
+                inputs=msg,
+                label="Example prompts for Shakti-100M",
+                fn=example_selector
+            )
+
+        with gr.Column(visible=False) as examples_250m:
+            gr.Examples(
+                examples=get_examples("Shakti-250M"),
+                inputs=msg,
+                label="Example prompts for Shakti-250M",
+                fn=example_selector
+            )
+
+        with gr.Column(visible=True) as examples_2_5b:
+            gr.Examples(
+                examples=get_examples("Shakti-2.5B"),
+                inputs=msg,
+                label="Example prompts for Shakti-2.5B",
+                fn=example_selector
+            )
+
+
+        # Update model selection and examples visibility
+        def combined_update(selected_model):
+            msg_update, chat_update = on_model_select(selected_model)
+            examples_100m_update, examples_250m_update, examples_2_5b_update = update_examples_visibility(
+                selected_model)
+            return [
+                msg_update,
+                chat_update,
+                examples_100m_update,
+                examples_250m_update,
+                examples_2_5b_update
+            ]
+
+
+        # Updated change event handler
+        model_dropdown.change(
+            combined_update,
+            inputs=[model_dropdown],
+            outputs=[
+                msg,
+                chatbot,
+                examples_100m,
+                examples_250m,
+                examples_2_5b
+            ]
+        )
 
-    demo.queue(max_size=20).launch()
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()