methods/ALGO/load_response.py

import json
import os


def read_jsonl(file_path):
    data = []
    with open(file_path, 'r') as file:
        for line in file:
            data.append(json.loads(line))
    return data

def read_pass_rate_file(file_path):
    data = {}
    with open(file_path, 'r') as file:
        for line in file:
            item = json.loads(line)
            data[item['tcb_id']] = {
                'gen_nums': item['gen_nums'],
                "right_nums": item['right_nums']
            }
    return data

import re

def extract_code(ans_str):
    pattern = r'```python\n(.*?)```'
    matches = re.findall(pattern, ans_str, re.DOTALL)
    return matches[-1]

def extract_content_code(ans_str):
    pattern = r'<ASSISTANT>(.*?)</ASSISTANT>'
    matches = re.findall(pattern, ans_str, re.DOTALL)
    return matches[-1]

def get_response_function(repsonse_path, model_name, test_al):   
    pass_rate_file = f"/home/luoxianzhen/yang/save_tests_{model_name}-new/algo/test_pass_rate.jsonl"
    pass_rate_dict = read_pass_rate_file(pass_rate_file)

    test_func_list = read_jsonl(repsonse_path.format(test_al, model_name))

    test_functions = {}
    
    for response_item in test_func_list:
        try:
            code_string = extract_code(response_item["response"])
        except:
            continue
        
        if response_item['tcb_id'] not in test_functions.keys():
            test_functions[response_item['tcb_id']] = {
                "random": [],
                "brute-force": []
            }
        test_functions[response_item['tcb_id']][response_item['type']].append(code_string)
    
    func_data = []
    for k, v in test_functions.items():
        if k not in pass_rate_dict.keys() or (pass_rate_dict[k]['gen_nums'] < 50 and pass_rate_dict[k]['right_nums'] < 10):
            func_data.append({
                "tcb_id": k,
                "func_list": v
            })
        # func_data.append({
        #         "tcb_id": k,
        #         "func_list": v
        #     })
    return func_data

def load_data(test_inputs):
    ds = json.load(open("/home/luoxianzhen/yang/data/Ours/TestcaseBench-v28.json", "r", encoding="utf-8"))
    res = []

    for item in ds:
        tests = [test_list for test_list in test_inputs if test_list["tcb_id"] == item['tcb_id']]
        if len(tests) <= 0:
            continue
        tests = tests[0]['generate_testcases']
        for c in (item['solutions'][0:3]):
            res.append({
                "code": c['code'],
                "time_limit": item["runtime_limit"],
                "memory_limit": item["memory_limit"],
                "compileAndRunOptions": c["compileAndRunOptions"],
                "test_cases": tests,
                "problem_id": item['tcb_id'],
            })
    return res
        
if __name__ == "__main__":
    get_response_function("/home/luoxianzhen/yang/data/response-orginal/orginal_response_{}_{}.jsonl", "qwen-coder-plus", 'algo')