from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import json
from tqdm import tqdm

model_path = ""
data_path=""
# default: Load the model on the available device(s)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    model_path, torch_dtype="auto", device_map="auto"
)


# default processor
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

# The default range for the number of visual tokens per image in the model is 4-16384.
# You can set min_pixels and max_pixels according to your needs, such as a token range of 256-1280, to balance performance and cost.
min_pixels = 256*28*28
max_pixels = 1003520
# processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)

with open(data_path, 'r', encoding='utf-8') as f:
         data = json.load(f)

for i, item in tqdm(list(enumerate(data)), total=len(data)):
    query = item['conversations'][0]['value']
    image_paths = item['image']
    converted_image_paths = []
    for path in image_paths:
        if path.startswith('box_images_new/'):
            new_path = path.replace(
                    'box_images_new/',
                    '/ex_disk2/kyunghwan/SlideVQA-main/box_images_new/'
                )
        # elif path.startswith('/ex_disk2/kyunghwan/'):
        #     new_path = path.replace(
        #          '/ex_disk2/kyunghwan/',
        #          '/ex_disk2/kyunghwan'
        #      )
        else:
            new_path = path
        converted_image_paths.append(new_path)

    segments = query.split("<image>")
    content_list = []



    for idx, seg in enumerate(segments):
        if seg.strip():
            content_list.append({"type": "text", "text": seg})

        if idx < len(converted_image_paths):
            content_list.append({"type": "image", "image": converted_image_paths[idx]})


    messages = [{"role": "user", "content": content_list}]
    

    # Preparation for inference
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to(model.device)

    # Inference: Generation of the output
    generated_ids = model.generate(**inputs, max_new_tokens=2000)
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    print(output_text)