Function batch_infer_local_chat_model

src/inference.py:869–949

(model, mode, rd, input_file, output_path)

def batch_infer_local_chat_model(model, mode, rd, input_file, output_path):
    if isinstance(model, str):
        model = LocalModel(model, "0", tensor_parallel_size = 8 if "70B" in model or "Mixtral" in model or "command" in model else 4)

    inputs = json.load(open(input_file, "r"))

    prompts = []
    overlong_prompts = []
    prompt_map = {}

    print("Generating prompts...")

    context_length = 8192 if "llama" in model.name else 128000

    for d in inputs:
        prompt = model.tokenizer.apply_chat_template(
            d["body"]["messages"],
            tokenize = False,
            add_generation_prompt=True
        )
        if model.get_prompt_length(prompt) < context_length and prompt not in prompt_map:
            prompts.append(prompt)
        elif prompt not in prompt_map:
            overlong_prompts.append(prompt)
        if prompt not in prompt_map:
            prompt_map[prompt] = []
        prompt_map[prompt].append(d["custom_id"])

    print("{} overlong prompts".format(len(overlong_prompts)))

    num = 0
    for prompt in prompt_map:
        num += len(prompt_map[prompt])

    if num != len(inputs):
        raise ValueError("The instances in the prompt map does not match the input size!")

    if len(prompt_map) != len(prompts) + len(overlong_prompts):
        raise ValueError("The number of prompt in the prompt map does not match the input prompts.")

        
    try:
        results = model.infer_many(prompts, n = inputs[0]["body"]["n"], temperature = inputs[0]["body"]["temperature"])
    except Exception as e:
        logger.error("Error occurred with reason: {}. INFO Dataset: {}\nModel: {}\n".format(str(e), dataset.name, model.name))
        traceback.print_exc()

    data = {}

    for prompt in results:
        for i in prompt_map[prompt]:
            if i.count("@") == 4:
                data[i] = {
                    "content": results[prompt][0][0],
                    "role": "assistant"
                }
            else:
                data[i] = {