|
from prompt import TA_prompt |
|
import re |
|
from utils import generate_response, run_code |
|
|
|
|
|
def post_process_code(code, question): |
|
func_name = code.split("(")[0].split("def")[-1].strip() |
|
parameters = code.split("\n")[0].split(f"def {func_name}")[-1][1:-2].split(",") |
|
if '' in parameters: |
|
parameters.remove('') |
|
values = re.findall(r"[-+]?\d*\.\d+|\d+", question)[:len(parameters)] |
|
values = [int(v) for v in values] |
|
arguments = list(zip(parameters, values)) |
|
|
|
arg_string = "" |
|
for param, val in arguments: |
|
arg_string += f"{param}={val}," |
|
func_call = f"\nprint({func_name}({arg_string[:-1]}))" |
|
code += func_call |
|
return code |
|
|
|
|
|
def solve_ta(question): |
|
question = question.strip() |
|
question = "Human: " + question |
|
query = TA_prompt + question |
|
query = query.strip() |
|
query += "\n" |
|
code = generate_response(query, 0.9) |
|
n = len(TA_prompt.strip()) |
|
code = code[n:].strip().split("-----")[0] |
|
|
|
splitting_string = "```" if "```python" not in code else "```python" |
|
if "```" in code: |
|
code = code.split(splitting_string)[1].split("```")[0].strip() |
|
|
|
code = post_process_code(code, question) |
|
print(code) |
|
|
|
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, re.DOTALL)[0].split("Human:")[0] |
|
return res.strip(), "" |
|
|
|
|
