Update app.py

app.py

@@ -3,17 +3,19 @@ import requests
 
 API_URL = "https://oocsd96vwjv6tw6a.us-east-1.aws.endpoints.huggingface.cloud"
 
-def query(message, history, system_message, temperature, top_p):
+def query(message, history, system_message, temperature, top_p, max_new_tokens, stop_tokens):
     if history is None:
         history = []
     full_prompt = f"{system_message}\n\n"
-    # Handle both new messages format and legacy tuple format
+    # Process conversation history, supporting both new and legacy formats.
     if history and isinstance(history[0], dict):
         for entry in history:
-            if entry.get("role") == "user":
-                full_prompt += f"User: {entry.get('content')}\n"
-            elif entry.get("role") == "assistant":
-                full_prompt += f"Assistant: {entry.get('content')}\n"
+            role = entry.get("role")
+            content = entry.get("content")
+            if role == "user":
+                full_prompt += f"User: {content}\n"
+            elif role == "assistant":
+                full_prompt += f"Assistant: {content}\n"
     elif history and isinstance(history, list):
         for pair in history:
             if isinstance(pair, (list, tuple)) and len(pair) == 2:
@@ -21,15 +23,20 @@ def query(message, history, system_message, temperature, top_p):
                 full_prompt += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
     full_prompt += f"User: {message}\nAssistant:"
     
+    # Process stop tokens input into a list if provided.
+    stop_list = [s.strip() for s in stop_tokens.split(",")] if stop_tokens.strip() != "" else None
+
     payload = {
         "inputs": full_prompt,
         "parameters": {
             "temperature": temperature,
             "top_p": top_p,
-            "max_new_tokens": 512,
-            "stop": ["User:", "Assistant:"]
+            "max_new_tokens": max_new_tokens,
         }
     }
+    if stop_list is not None:
+        payload["parameters"]["stop"] = stop_list
+
     headers = {"Content-Type": "application/json"}
     response = requests.post(f"{API_URL}/generate", json=payload, headers=headers)
     result = response.json()
@@ -45,14 +52,16 @@ demo = gr.ChatInterface(
     title="Monarch-1 Chatbot",
     description="Chat directly with Monarch-1, a generative AI optimized for African contexts.",
     examples=[
-        ["Tell me about ancient African civilizations.", None, "You are Monarch-1, an AI model optimized for African contexts.", 0.7, 0.9],
-        ["Explain economic trends in Africa.", None, "You are Monarch-1, an AI model optimized for African contexts.", 0.7, 0.9],
-        ["Teach me some Kiswahili phrases.", None, "You are Monarch-1, an AI model optimized for African contexts.", 0.7, 0.9]
+        ["Tell me about ancient African civilizations.", "You are Monarch-1, an AI model optimized for African contexts.", 0.7, 0.9, 512, "User:,Assistant:"],
+        ["Explain economic trends in Africa.", "You are Monarch-1, an AI model optimized for African contexts.", 0.7, 0.9, 512, "User:,Assistant:"],
+        ["Teach me some Kiswahili phrases.", "You are Monarch-1, an AI model optimized for African contexts.", 0.7, 0.9, 512, "User:,Assistant:"]
     ],
     additional_inputs=[
         gr.Textbox(value="You are Monarch-1, an AI model optimized for African contexts.", label="System message"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.5, maximum=1.0, value=0.9, step=0.05, label="Top-p (nucleus sampling)")
+        gr.Slider(minimum=0.5, maximum=1.0, value=0.9, step=0.05, label="Top-p (nucleus sampling)"),
+        gr.Slider(minimum=10, maximum=1024, value=512, step=10, label="Max New Tokens"),
+        gr.Textbox(value="User:,Assistant:", label="Stop Tokens (comma-separated)")
     ],
     type="messages"
 )