Update app.py

app.py

@@ -110,14 +110,22 @@ def chat_accordion():
 
 
 def format_chat_prompt(
-    message: str, chat_history, instructions: str, user_name: str, bot_name: str
+    message: str,
+    chat_history,
+    instructions: str,
+    user_name: str,
+    bot_name: str,
+    include_chat_history: bool = True,
 ) -> str:
     instructions = instructions.strip()
     prompt = instructions
-    for turn in chat_history:
-        user_message, bot_message = turn
-        prompt = f"{prompt}\n{user_name}: {user_message}\n{bot_name}: {bot_message}"
-    prompt = f"{prompt}\n{user_name}: {message}\n{bot_name}:"
+    if include_chat_history:
+        for turn in chat_history:
+            user_message, bot_message = turn
+            prompt = f"{prompt}\n{user_name}: {user_message}\n{bot_name}: {bot_message}"
+        prompt = f"{prompt}\n{user_name}: {message}\n{bot_name}:"
+    else:
+        prompt = f"{prompt}\n{user_name}: {message}\n{bot_name}:"
     return prompt
 
 
@@ -156,29 +164,26 @@ def chat_tab():
         session_id: str,
     ):
         prompt = format_chat_prompt(message, history, instructions, user_name, bot_name)
-        generated_response = ""
 
-        payload = json.dumps(
-            {
-                "endpoint": MODEL_NAME,
-                "data": {
-                    "inputs": prompt,
-                    "parameters": {
-                        "max_new_tokens": 1024,
-                        "do_sample": True,
-                        "top_p": top_p,
-                        "stop": ["User:"],
-                    },
-                    "stream": True,
-                    "session_id": session_id,
+        payload = {
+            "endpoint": MODEL_NAME,
+            "data": {
+                "inputs": prompt,
+                "parameters": {
+                    "max_new_tokens": 1024,
+                    "do_sample": True,
+                    "top_p": top_p,
+                    "stop": ["User:"],
                 },
-            }
-        )
+                "stream": True,
+                "session_id": session_id,
+            },
+        }
 
         sess = requests.Session()
         full_output = ""
         with sess.post(
-            ENDPOINT_URL, headers=HEADERS, data=payload, stream=True
+            ENDPOINT_URL, headers=HEADERS, json=payload, stream=True
         ) as response:
             if response.status_code == 200:
                 for chunk in response.iter_content(chunk_size=4):
@@ -191,7 +196,22 @@ def chat_tab():
                         else:
                             yield full_output
         if full_output == "":
-            yield "I am sorry, I did not understand your query. Could you please rephrase it?"
+            payload["data"]["inputs"] = format_chat_prompt(
+                message, history, instructions, user_name, bot_name, False
+            )
+            with sess.post(
+                ENDPOINT_URL, headers=HEADERS, json=payload, stream=True
+            ) as response:
+                if response.status_code == 200:
+                    for chunk in response.iter_content(chunk_size=4):
+                        if chunk:
+                            decoded = chunk.decode("utf-8")
+                            full_output += decoded
+                            if full_output.endswith("User:"):
+                                yield full_output[:-5]
+                                break
+                            else:
+                                yield full_output
         return ""
 
     with gr.Column():
@@ -268,4 +288,4 @@ def start_demo():
 
 
 if __name__ == "__main__":
-    start_demo()
+    start_demo()