Update app.py

app.py

@@ -24,7 +24,7 @@ except Exception as e:
     tokenizer = None
     peft_model = None
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
+async def respond(message, history, system_message, max_tokens, temperature, top_p):
     """
     Generates a response based on the user message and history using the provided PEFT model.
     Args:
@@ -56,7 +56,7 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
 
     # Generate the output
     try:
-        outputs = peft_model.generate(
+        outputs = await peft_model.generate(
             **inputs,
             max_new_tokens=max_tokens,
             temperature=temperature,
@@ -73,22 +73,23 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
     # Extract content between <user>...</user> tags
     def extract_user_content(text):
         """
-        Extracts and returns content between <user>...</user> tags or lines starting with 'output:' (case-insensitive) in the given text.
+        Extracts and returns content between <user>...</user> tags in the given text.
         If multiple such sections exist, their contents are concatenated.
         """
-        # Compile the regular expression pattern with re.IGNORECASE flag
-        pattern = re.compile(r'<user>(.*?)</user>|^output:\s*(.*)', re.IGNORECASE | re.DOTALL)
-        
-        # Find all matches in the text
-        matches = pattern.findall(text)
-        
-        # Extract and concatenate the matched content
-        extracted_content = '\n'.join(
-            (match[0] or match[1]).strip() for match in matches if match
-        )
-        
+        pattern = re.compile(r'<user>(.*?)</user>|output:', re.IGNORECASE)
+        matches = re.findall(pattern, text, re.DOTALL)
+        extracted_content = '\n'.join(match.strip() for match in matches)
         return extracted_content
 
+    # Extract the normalized text
+    normalized_text = extract_user_content(generated_text)
+
+    # Stream the response token by token
+    response = ""
+    for token in normalized_text.split():
+        response += token + " "
+        yield response.strip()
+
 # Gradio interface setup
 demo = gr.ChatInterface(
     respond,
@@ -105,4 +106,3 @@ demo = gr.ChatInterface(
 
 if __name__ == "__main__":
     demo.launch()
-