bigcode/Reasoning-with-StarCoder

1from prompt import TA_prompt2import re3from utils import generate_response, run_code4 5 6def post_process_code(code, question):7    func_name = code.split("(")[0].split("def")[-1].strip()8    parameters = code.split("\n")[0].split(f"def {func_name}")[-1][1:-2].split(",")9    if '' in parameters:10        parameters.remove('')11    values = re.findall(r"[-+]?\d*\.\d+|\d+", question)[:len(parameters)]12    values = [int(v) for v in values]13    arguments = list(zip(parameters, values))14 15    arg_string = ""16    for param, val in arguments:17        arg_string += f"{param}={val},"18    func_call = f"\nprint({func_name}({arg_string[:-1]}))"19    code += func_call20    return code21 22 23def solve_ta(question):24    question = question.strip()25    question = "Human: " + question26    query = TA_prompt + question27    query = query.strip()28    query += "\n"29    code = generate_response(query, 0.9)30    n = len(TA_prompt.strip())31    code = code[n:].strip().split("-----")[0]32    # print(code)33    splitting_string = "```" if "```python" not in code else "```python"34    if "```" in code:35        code = code.split(splitting_string)[1].split("```")[0].strip()36        # code preprocessing37        code = post_process_code(code, question)38        print(code)39        # code running40        if "input(" in code:41            return None, code42        pred = None43        try:44            pred = run_code(code)45        except Exception as ex:46            return None, code47        return pred, code48    else:49        res = re.findall(r"Assistant:(.*)", code, re.DOTALL)[0].split("Human:")[0]50        return res.strip(), ""51 52