rootacess commited on
Commit
950e174
1 Parent(s): 076bdf6

Performance improvement for TA + mathprompter

Browse files
Files changed (1) hide show
  1. TA.py +2 -2
TA.py CHANGED
@@ -46,8 +46,8 @@ def solve_ta(question, token):
46
  return None, code
47
  return pred, code
48
  else:
49
- res = re.findall(r"Assistant:(.*)", code)[0]
50
- return res, ""
51
 
52
 
53
  if __name__ == "__main__":
 
46
  return None, code
47
  return pred, code
48
  else:
49
+ res = re.findall(r"Assistant:(.*)", code, re.DOTALL)[0].split("Human:")[0]
50
+ return res.strip(), ""
51
 
52
 
53
  if __name__ == "__main__":