Update caption.py

caption.py

@@ -26,6 +26,7 @@ def process_file(llm, filepath, sampling_params):
     prev_eps = 0
     batch_texts = []
     df = pd.DataFrame()
+    batch_counter = 0  # Counter to keep track of batches processed
  
     if filepath.endswith('.parquet'):
         print("Reading from a parquet file...")
@@ -73,27 +74,30 @@ def process_file(llm, filepath, sampling_params):
                     executor.submit(df.to_parquet, filepath)
  
                 i += BATCH_SIZE
+                batch_counter += 1
+ 
+                # Push to hub every 10 batches
+                if batch_counter % 10 == 0:
+                    # Initialize the HuggingFace API
+                    api = HfApi()
+ 
+                    # Upload the processed file to the repository
+                    try:
+                        api.upload_file(
+                            path_or_fileobj=new_filepath,
+                            path_in_repo=new_filepath,
+                            repo_id="AlignmentLab-AI/caption_creation_0.8",
+                            repo_type="dataset",
+                        )
+                        print(f"Uploaded {new_filepath} to AlignmentLab-AI/caption_creation_0.8 repository.")
+                    except Exception as e:
+                        print(f"Error uploading file: {e}")
  
     # Delete the original parquet file if it is empty
     if df.empty:
         os.remove(filepath)
         print(f"Deleted the original file: {filepath}")
  
-        # Initialize the HuggingFace API
-        api = HfApi()
- 
-        # Upload the processed file to the repository
-        try:
-            api.upload_file(
-                path_or_fileobj=new_filepath,
-                path_in_repo=new_filepath,
-                repo_id="AlignmentLab-AI/caption_creation_0.8",
-                repo_type="dataset",
-            )
-            print(f"Uploaded {new_filepath} to AlignmentLab-AI/caption_creation_0.8 repository.")
-        except Exception as e:
-            print(f"Error uploading file: {e}")
- 
 def main():
     folder_name = 'captionate'
     sampling_params = SamplingParams(temperature=0.7, top_p=0.95, max_tokens=100)
@@ -108,4 +112,4 @@ def main():
  
 if __name__ == "__main__":
     main()
- 
+ `