Upload generate_prompt.py with huggingface_hub

generate_prompt.py

@@ -88,32 +88,31 @@ for batch_idx in tqdm(range(begin, end, batch_size)):
         data_list.append(messages)
         save_list.append(save_)
     #print(len(data_list))
-    text = processor.apply_chat_template(data_list, tokenize=False, add_generation_prompt=True)
-    #print(len(text))
-    image_inputs, video_inputs = process_vision_info(data_list)
-    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=128)
-    #print(generated_ids.shape)
-    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)
-    for idx,x in enumerate(output_text):
-        idx_real = batch_idx * batch_size + idx
-        save_list[idx][0]['result'] = x
-        save_data.append(save_list[idx])
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        #print(len(text))
+        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=128)
+        #print(generated_ids.shape)
+        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)
+        save_[0]['result'] = x
+        save_data.append(save_)
+
     if batch_idx % 4 ==0:
         write_json(json_path,save_data)
         print(len(save_data))