Update run.py

run.py

@@ -11,6 +11,13 @@ import hashlib
 import os
 import csv
 
+# pip install --no-cache-dir huggingface_hub[hf_transfer]
+def single_download(repo, fname, token, endpoint):
+    os.environ["TOKIO_WORKER_THREADS"] = "32"
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+    file = hf_hub_download(repo_id=repo, filename=fname, token=token, endpoint=endpoint, repo_type="dataset")
+    return file
+
 def load_hf_dataset(dataset_path, auth_token):
     dataset = load_dataset(dataset_path, token=auth_token)
     video_paths = dataset
@@ -29,27 +36,29 @@ def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, en
     
             if video_hf and video_hf_auth:
                 progress_info.append('Begin processing Hugging Face dataset.')
-                temp_parquet_file = hf_hub_download(
-                    repo_id=video_hf,
-                    filename='data/' + str(parquet_index).zfill(6) + '.parquet',
-                    repo_type="dataset",
-                    token=video_hf_auth,
-                )
-                parquet_file = pq.ParquetFile(temp_parquet_file)
-                for batch in parquet_file.iter_batches(batch_size=1):
+                os.environ["TOKIO_WORKER_THREADS"] = "8"
+                os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+                pqfile = hf_hub_download(repo_id=video_hf, filename='data/' + str(parquet_index).zfill(6) + '.parquet', repo_type="dataset", local_dir="/dev/shm", token=video_hf_auth)
+
+                pf = pq.ParquetFile(pqfile)
+                for batch in pf.iter_batches(1):
+                    _chunk = []
                     df = batch.to_pandas()
-                    video = df['video'][0]
-                    md5 = hashlib.md5(video).hexdigest()
-                    with tempfile.NamedTemporaryFile(dir=temp_dir) as temp_file:
-                        temp_file.write(video)
-                        video_path = temp_file.name
-                        processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
-                        frames = processor._decode(video_path)
-                        base64_list = processor.to_base64_list(frames)
-                        api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
-                        caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
-                        writer.writerow({'md5': md5, 'caption': caption})
-                        progress_info.append(f"Processed video with MD5: {md5}")
+                    for binary in df["video"]:
+                        if(binary):
+                            _v = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+                            with open(_v.name, "wb") as f:
+                                _ = f.write(binary)
+                            _chunk.append(_v.name)
+                            
+                            processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
+                            frames = processor._decode(_v.name)
+                            base64_list = processor.to_base64_list(frames)
+                            api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
+                            caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
+                            writer.writerow({'md5': _v.name, 'caption': caption})
+                            progress_info.append(f"Processed video with MD5: {md5}")
+                    return csv_filename, "\n".join(progress_info), None
                 return csv_filename, "\n".join(progress_info), None
             else:
                 return "", "No video source selected.", None