Update run.py

run.py

@@ -49,6 +49,7 @@ def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, en
                 )
 
                 pf = pq.ParquetFile(pqfile)
+                i = 0
                 for batch in pf.iter_batches(1):
                     _chunk = []
                     df = batch.to_pandas()
@@ -68,6 +69,9 @@ 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}")
+                            i += 1
+                        if i == 10:
+                            break
                         #return csv_filename, "\n".join(progress_info), None
                 return csv_filename, "\n".join(progress_info), None
             else: