Update app.py

app.py

@@ -29,9 +29,11 @@ model_options = {
 # Initialize tokenizer and model variables
 tokenizer = None
 model = None
+current_model = "Shakti-2.5B"  # Keep track of current model
+
 
 def load_model(selected_model: str):
-    global tokenizer, model
+    global tokenizer, model, current_model
     model_id = model_options[selected_model]
     tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.getenv("SHAKTI"))
     model = AutoModelForCausalLM.from_pretrained(
@@ -41,29 +43,44 @@ def load_model(selected_model: str):
         token=os.getenv("SHAKTI")
     )
     model.eval()
+    current_model = selected_model  # Update the current model
+
 
 # Initial model load (default to 2.5B)
 load_model("Shakti-2.5B")
 
+
 @spaces.GPU(duration=90)
 def generate(
-    message: str,
-    chat_history: list[tuple[str, str]],
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
+        message: str,
+        chat_history: list[tuple[str, str]],
+        max_new_tokens: int = 1024,
+        temperature: float = 0.6,
+        top_p: float = 0.9,
+        top_k: int = 50,
+        repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
     conversation = []
-    for user, assistant in chat_history:
-        conversation.extend(
-            [
-                json.loads(os.getenv("PROMPT")),
-                {"role": "user", "content": user},
-                {"role": "assistant", "content": assistant},
-            ]
-        )
+
+    # 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},
+                ]
+            )
+    else:
+        for user, assistant in chat_history:
+            conversation.extend(
+                [
+                    {"role": "user", "content": user},
+                    {"role": "assistant", "content": assistant},
+                ]
+            )
+
     conversation.append({"role": "user", "content": message})
 
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
@@ -92,51 +109,34 @@ def generate(
         outputs.append(text)
         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"]]
+                ["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?"]]
+                ["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"], ['मुझे भारतीय इतिहास के बारे में बताएं']]
+        return [["Tell me a story"], ["write a short poem which is hard to sing"],
+                ['मुझे भारतीय इतिहास के बारे में बताएं']]
+
 
 def on_model_select(selected_model):
     load_model(selected_model)  # Load the selected model
-    return update_examples(selected_model)  # Return new examples based on the selected model
-
-
-chat_interface = gr.ChatInterface(
-    fn=generate,
-    additional_inputs=[
-        gr.Slider(
-            label="Max new tokens",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.1,
-            maximum=4.0,
-            step=0.1,
-            value=0.6,
-        ),
-    ],
-    stop_btn=None,
-    examples=update_examples("Shakti-2.5B"),  # Set initial examples for 2.5B model
-    cache_examples=False,
-)
+    examples = update_examples(selected_model)  # Update examples
+    return gr.update(examples=examples), gr.update(value=[])  # Clear the chat space and update examples
+
+
+chat_history = gr.Chatbot()
 
 with gr.Blocks(css="style.css", fill_height=True) as demo:
     gr.Markdown(DESCRIPTION)
@@ -150,10 +150,32 @@ with gr.Blocks(css="style.css", fill_height=True) as demo:
         interactive=True,
     )
 
-    # Function to handle model change and update examples dynamically
-    model_dropdown.change(on_model_select, inputs=model_dropdown, outputs=[chat_interface])
+    # 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,
+    )
+
+    temperature_slider = gr.Slider(
+        label="Temperature",
+        minimum=0.1,
+        maximum=4.0,
+        step=0.1,
+        value=0.6,
+    )
+
+    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,
+    )
 
-    chat_interface.render()
+    # Function to handle model change and update examples dynamically
+    model_dropdown.change(on_model_select, inputs=model_dropdown, outputs=[chat_interface, chat_history])
 
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.launch()