Update app.py

app.py

@@ -17,6 +17,8 @@ except Exception as e:
     tokenizer = None
     base_model = None
     peft_model = None
+    
+generated_text = None
 
 def respond(message, history, system_message, max_tokens, temperature, top_p):
     global tokenizer, peft_model
@@ -47,20 +49,24 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         
         # Get the last message from the generated text
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        def extract_user_content(text):
-            """
-            Extracts and returns content between <user>...</user> tags in the given text.
-            If multiple such sections exist, their contents are concatenated.
-            """
-            # Compile the pattern with the desired flags
-            pattern = re.compile(r'<user>(.*?)</user>', re.IGNORECASE | re.DOTALL)
-            matches = re.findall(pattern, text)
-            extracted_content = '\n'.join(match.strip() for match in matches)
-            return extracted_content
+        def extract_content_after_output(text):
+    """
+    Extracts and returns the content that follows the word 'output' in the given text.
+    If 'output' is not found, returns an empty string.
+    """
+    # Define the regular expression pattern to match 'output' and capture everything after it
+    pattern = re.compile(r'output\s*:\s*(.*)', re.IGNORECASE | re.DOTALL)
+    match = pattern.search(text)
+    if match:
+        # Return the captured group, which is the content after 'output'
+        return match.group(1).strip()
+    else:
+        # If 'output' is not found, return an empty string
+        return "Retry to get output, the model failed to generated required output(This occurs rarely🤷‍♂️)"
         
-        # lines = extract_user_content(generated_text)
+        lines = extract_user_content(generated_text)
         
-        return generated_text
+        return lines
         
     except Exception as e:
         return f"Generation error: {e}"