Function single_process

preprocess/caption_qwen_multi.py:29–242

(input_csv_folder_path, store_csv_folder_path, GPU_offset)

def single_process(input_csv_folder_path, store_csv_folder_path, GPU_offset):
    

    # Setting
    store_freq = 10
    device = 'cuda'



    # Read the csv file
    csv_idx = GPU_offset
    csv_file_path = os.path.join(input_csv_folder_path, "sub" + str(csv_idx) + ".csv")
    print("CSV file we read is ", csv_file_path)


    # Prepare the store file path
    store_file_path = os.path.join(store_csv_folder_path, "sub" + str(csv_idx) + ".csv")
    if not resume and os.path.exists(store_file_path):
        # Remove existing csv
        os.remove(store_file_path)

    # Resume the store csv  
    find_resume = True
    if resume:      # Read the last store row
        find_resume = False
        with open(store_file_path, 'r') as file:
            reader = csv.reader(file)
            store_rows = list(reader)
            last_store_row = store_rows[-1]
            print("The number of rows we have processed in the store csv is ", len(store_rows))
            

    # Init the model
    processor = AutoProcessor.from_pretrained(model_path)
    # bnb_config = BitsAndBytesConfig(
    #                                     load_in_4bit=True,
    #                                     bnb_4bit_compute_dtype=torch.float16,  # Use float16 for computations
    #                                     bnb_4bit_use_double_quant=True,
    #                                     bnb_4bit_quant_type='nf4',  # NormalFloat4 quantization
    #                                 )
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                                                                    model_path, torch_dtype="auto", device_map="auto",
                                                                    # quantization_config=bnb_config,
                                                                )
    

    # Read all row in the csv file
    start_time = time.time()
    info_lists = []       # The order will be follow automatically
    with open(csv_file_path) as file_obj:
        reader_obj = csv.reader(file_obj) 
        
        # Iterate over each row in the csv  
        for row_idx, row in enumerate(reader_obj): 

            # For the first row case (With all title content)
            if row_idx == 0:    # The first line is the title of content
                elements = dict()