Break up Huggingface Upload

src/core/tasks/scheduled/impl/huggingface/operator.py

@@ -1,3 +1,4 @@
+from itertools import count
 
 from src.core.tasks.scheduled.templates.operator import ScheduledTaskOperatorBase
 from src.db.client.async_ import AsyncDatabaseClient
@@ -30,7 +31,10 @@
 
         # Otherwise, push to huggingface
         run_dt = await self.adb_client.get_current_database_time()
-        outputs = await self.adb_client.get_data_sources_raw_for_huggingface()
-        self.hf_client.push_data_sources_raw_to_hub(outputs)
+        for idx in count(start=1):
+            outputs = await self.adb_client.get_data_sources_raw_for_huggingface(page=idx)
+            if len(outputs) == 0:
+                break
+            self.hf_client.push_data_sources_raw_to_hub(outputs, idx=idx)
 
         await self.adb_client.set_hugging_face_upload_state(run_dt.replace(tzinfo=None))

src/core/tasks/scheduled/impl/huggingface/queries/get/core.py

@@ -5,6 +5,7 @@
 from src.core.tasks.scheduled.impl.huggingface.queries.get.convert import convert_url_status_to_relevant, \
     convert_fine_to_coarse_record_type
 from src.core.tasks.scheduled.impl.huggingface.queries.get.model import GetForLoadingToHuggingFaceOutput
+from src.db.client.helpers import add_standard_limit_and_offset
 from src.db.models.impl.url.html.compressed.sqlalchemy import URLCompressedHTML
 from src.db.models.impl.url.core.sqlalchemy import URL
 from src.db.queries.base.builder import QueryBuilderBase
@@ -13,6 +14,10 @@
 
 class GetForLoadingToHuggingFaceQueryBuilder(QueryBuilderBase):
 
+    def __init__(self, page: int):
+        super().__init__()
+        self.page = page
+
 
     async def run(self, session: AsyncSession) -> list[GetForLoadingToHuggingFaceOutput]:
         label_url_id = 'url_id'
@@ -42,6 +47,7 @@
                 ])
             )
         )
+        query = add_standard_limit_and_offset(page=self.page, statement=query)
         db_results = await sh.mappings(
             session=session,
             query=query

src/db/client/async_.py

@@ -1463,9 +1463,9 @@
             )
             session.add(compressed_html)
 
-    async def get_data_sources_raw_for_huggingface(self) -> list[GetForLoadingToHuggingFaceOutput]:
+    async def get_data_sources_raw_for_huggingface(self, page: int) -> list[GetForLoadingToHuggingFaceOutput]:
         return await self.run_query_builder(
-            GetForLoadingToHuggingFaceQueryBuilder()
+            GetForLoadingToHuggingFaceQueryBuilder(page)
         )
 
     async def set_hugging_face_upload_state(self, dt: datetime) -> None:

src/external/huggingface/hub/client.py

@@ -1,5 +1,6 @@
 
 from datasets import Dataset
+from huggingface_hub import HfApi
 
 from src.external.huggingface.hub.constants import DATA_SOURCES_RAW_REPO_ID
 from src.external.huggingface.hub.format import format_as_huggingface_dataset
@@ -10,22 +11,39 @@
 
     def __init__(self, token: str):
         self.token = token
+        self.api = HfApi(token=token)
 
-    def _push_dataset_to_hub(self, repo_id: str, dataset: Dataset) -> None:
+    def _push_dataset_to_hub(
+        self,
+        repo_id: str,
+        dataset: Dataset,
+        idx: int
+    ) -> None:
         """
         Modifies:
             - repository on Hugging Face, identified by `repo_id`
         """
-        dataset.push_to_hub(repo_id=repo_id, token=self.token)
+        dataset.to_parquet(f"part_{idx}.parquet")
+        self.api.upload_file(
+            path_or_fileobj=f"part_{idx}.parquet",
+            path_in_repo=f"data/part_{idx}.parquet",
+            repo_id=repo_id,
+            repo_type="dataset",
+        )
 
     def push_data_sources_raw_to_hub(
         self,
-        outputs: list[GetForLoadingToHuggingFaceOutput]
+        outputs: list[GetForLoadingToHuggingFaceOutput],
+        idx: int
     ) -> None:
         """
         Modifies:
             - repository on Hugging Face, identified by `DATA_SOURCES_RAW_REPO_ID`
         """
         dataset = format_as_huggingface_dataset(outputs)
         print(dataset)
-        self._push_dataset_to_hub(repo_id=DATA_SOURCES_RAW_REPO_ID, dataset=dataset)
+        self._push_dataset_to_hub(
+            repo_id=DATA_SOURCES_RAW_REPO_ID,
+            dataset=dataset,
+            idx=idx
+        )