merve HF staff commited on
Commit
458ccb5
1 Parent(s): a3a174a

Fix regex that yield conv turns

Browse files
Files changed (1) hide show
  1. app.py +8 -6
app.py CHANGED
@@ -19,12 +19,14 @@ pipe = pipeline("image-to-text", model=model_id, model_kwargs={"quantization_con
19
 
20
 
21
  def extract_response_pairs(text):
22
- pattern = re.compile(r'(USER:.*?)ASSISTANT:(.*?)(?:$|USER:)', re.DOTALL)
23
- matches = pattern.findall(text)
24
-
25
- pairs = [(user.strip(), assistant.strip()) for user, assistant in matches]
26
-
27
- return pairs
 
 
28
 
29
 
30
  def postprocess_output(output: str) -> str:
 
19
 
20
 
21
  def extract_response_pairs(text):
22
+ turns = re.split(r'(USER:|ASSISTANT:)', text)[1:]
23
+ turns = [turn.strip() for turn in turns if turn.strip()]
24
+ print(turns[1::2])
25
+ conv_list = []
26
+ for i in range(0, len(turns[1::2]), 2):
27
+ if i + 1 < len(turns[1::2]):
28
+ conv_list.append((turns[1::2][i].lstrip(":"), turns[1::2][i + 1].lstrip(":")))
29
+ return conv_list
30
 
31
 
32
  def postprocess_output(output: str) -> str: