Upload generate_prompt.py with huggingface_hub

generate_prompt.py

@@ -18,32 +18,42 @@ def write_json(file_path, data):
 
 # default: Load the model on the available device(s)
 print(torch.cuda.device_count())
-model_path = "/home/zbz5349/WorkSpace/aigeeks/Qwen2.5-VL/ckpt_7B"
+#model_path = "/home/zbz5349/WorkSpace/aigeeks/Qwen2.5-VL/ckpt_7B"
 # model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 #     model_path, torch_dtype="auto", device_map="auto"
 # )
 
+
 # We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_path", type=str, default="/home/zbz5349/WorkSpace/aigeeks/Qwen2.5-VL/ckpt_7B")
+parser.add_argument("--begin", type=int, default=0)
+parser.add_argument("--end", type=int, default=4635)
+parser.add_argument("--batch_size", type=int, default=3)
+parser.add_argument("--data_path", type=str, default="/home/zbz5349/WorkSpace/aigeeks/Qwen2.5-VL/magicbrush_dataset/dataset.json")
+parser.add_argument("--prompt_path", type=str, default="/home/zbz5349/WorkSpace/aigeeks/Qwen2.5-VL/magicbrush_dataset/gen.json")
+
+args = parser.parse_args()
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    model_path,
+    agrs.model_path,
     torch_dtype=torch.bfloat16,
     attn_implementation="flash_attention_2",
     device_map="auto",
 )
 
 # default processor
-processor = AutoProcessor.from_pretrained(model_path)
+processor = AutoProcessor.from_pretrained(args.model_path)
 print(model.device)
 
 
 
 
-data = read_json('/home/zbz5349/WorkSpace/aigeeks/Qwen2.5-VL/magicbrush_dataset/dataset.json')
+data = read_json(args.data_path)
 save_data = []
 correct_num = 0
-begin = 0
-end = len(data)
-batch_size = 1
+begin = args.begin
+end = args.end
+batch_size = args.batch_size
 for batch_idx in tqdm(range(begin, end, batch_size)):
     batch = data[batch_idx:batch_idx + batch_size] 
     
@@ -99,7 +109,7 @@ for batch_idx in tqdm(range(begin, end, batch_size)):
         save_list[idx]['result'] = x
         save_data.append(save_list[idx])
 
-json_path = "image_path = '/home/zbz5349/WorkSpace/aigeeks/Qwen2.5-VL/magicbrush_dataset/gen.json'"
+json_path = args.prompt_path
 write_json(json_path,save_data)