Update app.py

app.py

@@ -84,12 +84,19 @@ model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map='balanced')
 def filter_text(text):
   """Filters out the text up to and including 'Rewritten Prompt:'."""
   phrase = "Rewritten Prompt: "
+  phraseB = "rewritten text: "
   pattern = f"(.*?){re.escape(phrase)}(.*)"
-  match = re.search(pattern, text)
+  patternB = f"(.*?){re.escape(phraseB)}(.*)"
+  matchB = re.search(patternB, text)
  # match = re.search(pattern, text,flags=re.DOTALL)
-  if match:
-        filtered_text = match.group(2)
-        return filtered_text
+  if matchB:
+        filtered_text = matchB.group(2)
+        match = re.search(pattern, filtered_text)
+        if match:
+          filtered_text = match.group(2)
+          return filtered_text
+        else:
+          return filtered_text
   else:
         # Handle the case where no match is found
         return text