Function run_inference

tools/caption/caption_llava_next.py:226–370

Run inference on ActivityNet QA DataSet using the Video-ChatGPT model. Args: args: Command-line arguments.

(rank, world_size, args)

def run_inference(rank, world_size, args):
    """
    Run inference on ActivityNet QA DataSet using the Video-ChatGPT model.

    Args:
        args: Command-line arguments.
    """
    setup(rank, world_size)

    device = torch.device(f"cuda:{rank}")
    # Initialize the model
    model_name = get_model_name_from_path(args.model_path)
    # Set model configuration parameters if they exist
    if args.overwrite == True:
        overwrite_config = {}
        overwrite_config["mm_spatial_pool_mode"] = args.mm_spatial_pool_mode
        overwrite_config["mm_spatial_pool_stride"] = args.mm_spatial_pool_stride
        overwrite_config["mm_newline_position"] = args.mm_newline_position

        cfg_pretrained = AutoConfig.from_pretrained(args.model_path)

        # import pdb;pdb.set_trace()
        if "qwen" not in args.model_path.lower():
            if "224" in cfg_pretrained.mm_vision_tower:
                # suppose the length of text tokens is around 1000, from bo's report
                least_token_number = args.for_get_frames_num * (16 // args.mm_spatial_pool_stride) ** 2 + 1000
            else:
                least_token_number = args.for_get_frames_num * (24 // args.mm_spatial_pool_stride) ** 2 + 1000

            scaling_factor = math.ceil(least_token_number / 4096)
            if scaling_factor >= 2:
                if "vicuna" in cfg_pretrained._name_or_path.lower():
                    print(float(scaling_factor))
                    overwrite_config["rope_scaling"] = {"factor": float(scaling_factor), "type": "linear"}
                overwrite_config["max_sequence_length"] = 4096 * scaling_factor
                overwrite_config["tokenizer_model_max_length"] = 4096 * scaling_factor
        if args.load_8bit:
            quantization_config = BitsAndBytesConfig(load_in_8bit=True, bnb_8bit_compute_dtype=torch.bfloat16)
            tokenizer, model, image_processor, context_len = load_pretrained_model(
                args.model_path,
                args.model_base,
                model_name,
                device_map=device,
                quantization_config=quantization_config,
                overwrite_config=overwrite_config,
            )
        else:
            tokenizer, model, image_processor, context_len = load_pretrained_model(
                args.model_path, args.model_base, model_name, device_map=device, overwrite_config=overwrite_config
            )
    else:
        tokenizer, model, image_processor, context_len = load_pretrained_model(
            args.model_path, args.model_base, model_name, device_map=device
        )

    if tokenizer.pad_token_id is None:
        if "qwen" in tokenizer.name_or_path.lower():
            # print("Setting pad token to bos token for qwen model.")