Performance improvement for TA + mathprompter

TA.py

@@ -27,27 +27,32 @@ def solve_ta(question, token):
     query = query.strip()
     query += "\n"
     code = generate_response(query, 0.9, token)
-
+    n = len(TA_prompt.strip())
+    code = code[n:].strip().split("-----")[0]
+    # print(code)
     splitting_string = "```" if "```python" not in code else "```python"
-    code = code.split(splitting_string)[-2].split("```")[0].strip()
-    print(code)
-    # code preprocessing
-
-    code = post_process_code(code, question)
-    print(code)
-    # code running
-    if "input(" in code:
-        return None, code
-    pred = None
-    try:
-        pred = run_code(code)
-    except Exception as ex:
-        return None, code
-    return pred, code
+    if "```" in code:
+        code = code.split(splitting_string)[1].split("```")[0].strip()
+        # code preprocessing
+        code = post_process_code(code, question)
+        print(code)
+        # code running
+        if "input(" in code:
+            return None, code
+        pred = None
+        try:
+            pred = run_code(code)
+        except Exception as ex:
+            return None, code
+        return pred, code
+    else:
+        res = re.findall(r"Assistant:(.*)", code)[0]
+        return res, ""
 
 
 if __name__ == "__main__":
 
-    q = "What is the 7th Fibonacci number? Write the Python code"
-    print(solve_ta(q))
+    q = "What is the smallest even prime number?"
+    # q = "Write the code to find the 1st even prime number"
+    print(solve_ta(q, "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"))
 

app.py

@@ -42,6 +42,8 @@ def run(token, question, method):
 
 
 def run_edits(code):
+    if "input(" in code:
+        return "Code execution failed, Please remove any input statement or bugs", code
     try:
         code_op = run_code(code)
         return code_op, code
@@ -90,10 +92,10 @@ with demo:
 
     access_token = gr.Textbox(type="password", label="Hugging Face Access Token")
     with gr.Row():
-        methods = gr.Dropdown(choices=['PAL', 'TA', 'MathPrompter'], interactive=True, label="Evaluation Strategies")
+        methods = gr.Dropdown(choices=['PAL', 'TA', 'MathPrompter'], value="PAL",interactive=True, label="Evaluation Strategies")
         question_input = gr.Textbox(label="Question", lines=1, placeholder="Enter your question here...")
 
-    instruction = gr.Textbox(label="Instructions", visible=False, interactive=False)
+    instruction = gr.Textbox(label="Instructions", visible=True, interactive=False, value=render_instruction("PAL")['value'])
     methods.change(fn=render_instruction, inputs=methods, outputs=instruction)
 
     question_output = gr.Textbox(label="Answer", interactive=False)
@@ -123,3 +125,9 @@ with demo:
     )
 
 demo.launch()
+
+'''
+Carol was playing a trivia game. In the first round she scored 17 points and in the second round she scored 6 points. In the last round she lost 16 points. How many points did she have at the end of the game?
+While on vacation, Debby took 24 pictures at the zoo and 12 at the museum. If she later deleted 14 of the pictures, how many pictures from her vacation did she still have?
+I had 5 apples, I gave 3 to my brother and 4 to my uncle. How many apples I am left with? 
+'''

mathprompter.py

@@ -19,7 +19,7 @@ def generate_algebric_template(question):
     return question, var_map
 
 
-def generate_algebric_expression(question, variables, param, token):
+def generate_algebric_expression(question, param, token):
     question = question.strip()
     query = algebric_prompt.format(question=question).strip() + "\n"
     response = generate_response(query, param, token)
@@ -36,6 +36,7 @@ def generate_python_code(question, equation, param, token):
 
 def run(question, random_candidates, hps, token):
     question, var_map = generate_algebric_template(question)
+    print(question)
 
     # generating the random candidates for arguments
     random_mapping = pd.DataFrame(columns=list(var_map.keys()))
@@ -49,8 +50,7 @@ def run(question, random_candidates, hps, token):
     N = len(hps)
     for i in range(N):
 
-        variables = list(var_map.keys())
-        expression = generate_algebric_expression(question, variables, hps[i], token)
+        expression = generate_algebric_expression(question, hps[i], token)
         code = generate_python_code(question, expression, hps[i], token)
         candidates.append((expression, code))
         current_acc = 0

prompt.py

@@ -1,5 +1,3 @@
-# todo: Improve the prompt of mathprompter too
-
 prompt = '''
 def solution():
     #Ques: For Halloween Debby and her sister combined the candy they received. Debby had 32 pieces of candy while her sister had 42. If they ate 35 pieces the first night, how many pieces do they have left?
@@ -58,10 +56,6 @@ algebric_prompt = '''
 Answer = (A + B) - C
 
 
-#Ques: Roger had A dollars. If he spent B bucks on a new game, how many C dollar toys could he buy with the money he had left?
-Answer = (A - B) // C
-
-
 #Ques: A waiter had A customers in his section. If B of them left and the rest of his tables had C people at each table, how many tables did he have?
 Answer = (A - B) / C
 
@@ -75,15 +69,15 @@ Answer = (A - B) * C
 
 
 python_prompt = '''
-#Ques: For Halloween Debby and her sister combined the candy they received. Debby had A pieces of candy while her sister had B. If they ate C pieces the first night, how many pieces do they have left?
+# Ques: For Halloween Debby and her sister combined the candy they received. Debby had A pieces of candy while her sister had B . If they ate C pieces the first night, how many pieces do they have left?
 # Expression = (A + B) - C
 
-# Function for above Expression is:
+# Function for above Expression and Ques is:
 def solution(A, B, C):
     return (A + B) - C
 
 
-#Ques: A waiter had A customers in his section. If B of them left and the rest of his tables had C people at each table, how many tables did he have?    
+# Ques: A waiter had A customers in his section. If B of them left and the rest of his tables had C people at each table, how many tables did he have?    
 # Expression = (A - B) / C
 
 # Function for above expression is:
@@ -91,15 +85,7 @@ def solution(A, B, C):
     return (A - B) / C
 
 
-#Ques: Roger had A dollars. If he spent B bucks on a new game, how many C dollar toys could he buy with the money he had left?
-# Expression = (A - B) // C
-
-# Function for above expression is:
-def solution(A, B, C):
-    return (A - B) // C
-
-
-#Ques: There were A friends playing a video game online when B players quit. If each player left had C lives, how many lives did they have total? "
+# Ques: There were A friends playing a video game online when B more players joined the game. If each player had C lives, how many lives did they have total?
 # Expression = (A - B) * C
 
 # Function for above expression is: