Function batch_infer_local_model

src/inference.py:829–866

(model, dataset, output_path, device_num, n = 200, temperature = 0.7, dataset_repo = "test_dataset", prefix = None, suffix = None, chat = False, context_length = None)

def batch_infer_local_model(model, dataset, output_path, device_num, n = 200, temperature = 0.7, dataset_repo = "test_dataset", prefix = None, suffix = None, chat = False, context_length = None):
    if isinstance(model, str):
        model = LocalModel(model, device_num)
    
    if dataset != "mbpp":
        dataset = Dataset(dataset, data_path = os.path.join(dataset_repo, dataset.replace("/", "_")), full = False)
    else:
        dataset = Dataset(dataset, data_path = os.path.join(dataset_repo, dataset.replace("/", "_")), testfile_path = "datasets/MbppPlus.jsonl", full = False)

    if not chat:
        prompts, overlong_prompts = dataset.get_all_prompts(model = model, context_length = context_length)
        if prefix != None or suffix != None:
            new_prompts = []
            for prompt in prompts:
                new_prompts.append((prefix if prefix != None else "") + prompt + (suffix if suffix != None else ""))
            prompts = new_prompts
        print("{} overlong prompts found.".format(len(overlong_prompts)))
    else:
        dialogs = dataset.get_all_chats()
    try:
        if not chat:
            results = model.infer_many(prompts, n = n, temperature = temperature)
            for p in overlong_prompts:
                results[p] = ["Overlong input prompts.", False]
        else:
            results = model.infer_many_chats(dialogs)
    except Exception as e:
        logger.error("Error occurred with reason: {}. INFO Dataset: {}\nModel: {}\n".format(str(e), dataset.name, model.name))
        traceback.print_exc()

    res_dir_path = os.path.join(output_path, dataset.name.replace("/", "_"))
    res_path = os.path.join(res_dir_path, model.name.replace("/", "_") + ".json")

    if not os.path.exists(res_dir_path):
        os.mkdir(res_dir_path)
    
    with open(res_path, "w", encoding="utf-8") as of:
        of.write(json.dumps(results, sort_keys=True, indent=4, separators=(',', ': ')))


def batch_infer_local_chat_model(model, mode, rd, input_file, output_path):