Update run.py

run.py

@@ -28,6 +28,7 @@ def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, en
     progress_info = []
     processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
     api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
+    ind = 0
     with tempfile.TemporaryDirectory() as temp_dir:
         # temp_dir = '/opt/run'
         csv_filename = os.path.join('/dev/shm', str(parquet_index).zfill(6) + '_gpt4o_caption.csv')
@@ -54,6 +55,7 @@ def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, en
                     _chunk = []
                     df = batch.to_pandas()
                     for binary in df["video"]:
+                        ind += 1
                         if(binary):
                             _v = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
                             with open(_v.name, "wb") as f:
@@ -67,7 +69,8 @@ def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, en
                             writer.writerow({'md5': md5, 'caption': caption})
                             # writer.writerow({'md5': md5, 'caption': 'caption'})
                             progress_info.append(f"Processed video with MD5: {md5}")
-                        # return csv_filename, "\n".join(progress_info), None
+                        if ind == 2:
+                            return csv_filename, "\n".join(progress_info), None
                 return csv_filename, "\n".join(progress_info), None
             else:
                 return "", "No video source selected.", None