Fix regex that yield conv turns

app.py

@@ -19,12 +19,14 @@ pipe = pipeline("image-to-text", model=model_id, model_kwargs={"quantization_con
 
 
 def extract_response_pairs(text):
-    pattern = re.compile(r'(USER:.*?)ASSISTANT:(.*?)(?:$|USER:)', re.DOTALL)
-    matches = pattern.findall(text)
-
-    pairs = [(user.strip(), assistant.strip()) for user, assistant in matches]
-
-    return pairs
+    turns = re.split(r'(USER:|ASSISTANT:)', text)[1:]
+    turns = [turn.strip() for turn in turns if turn.strip()]
+    print(turns[1::2])
+    conv_list = []
+    for i in range(0, len(turns[1::2]), 2):
+        if i + 1 < len(turns[1::2]):
+            conv_list.append((turns[1::2][i].lstrip(":"), turns[1::2][i + 1].lstrip(":")))
+    return conv_list
 
 
 def postprocess_output(output: str) -> str: