nucliadb-dataset 6.4.2.post4389__py3-none-any.whl → 6.5.0.post4404__py3-none-any.whl

nucliadb_dataset/dataset.py

@@ -15,7 +15,7 @@
 import logging
 import os
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 
 import pyarrow as pa  # type: ignore
 
@@ -29,8 +29,9 @@ from nucliadb_dataset.tasks import (
 from nucliadb_models.entities import KnowledgeBoxEntities
 from nucliadb_models.labels import KnowledgeBoxLabels
 from nucliadb_models.resource import KnowledgeBoxObj
+from nucliadb_models.trainset import TrainSet as TrainSetModel
 from nucliadb_models.trainset import TrainSetPartitions
-from nucliadb_protos.dataset_pb2 import TrainSet
+from nucliadb_protos.dataset_pb2 import TrainSet as TrainSetPB
 from nucliadb_sdk.v2.sdk import NucliaDB
 
 logger = logging.getLogger("nucliadb_dataset")
@@ -104,7 +105,7 @@ class NucliaDBDataset(NucliaDataset):
         kbid: str,
         task: Optional[Task] = None,
         labels: Optional[List[str]] = None,
-        trainset: Optional[TrainSet] = None,
+        trainset: Optional[Union[TrainSetPB, TrainSetModel]] = None,
         base_path: Optional[str] = None,
         search_sdk: Optional[NucliaDB] = None,
         reader_sdk: Optional[NucliaDB] = None,
@@ -119,7 +120,7 @@ class NucliaDBDataset(NucliaDataset):
             task_definition = TASK_DEFINITIONS.get(task)
             if task_definition is None:
                 raise KeyError("Not a valid task")
-            trainset = TrainSet(type=task_definition.proto)
+            trainset = TrainSetPB(type=task_definition.proto)
             if task_definition.labels:
                 trainset.filter.labels.extend(labels)
         elif trainset is not None:

nucliadb_dataset/streamer.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 
 import logging
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 
 import requests
 
-from nucliadb_protos.dataset_pb2 import TrainSet
+from nucliadb_models.trainset import TrainSet as TrainSetModel
+from nucliadb_protos.dataset_pb2 import TrainSet as TrainSetPB
 
 logger = logging.getLogger("nucliadb_dataset")
 
@@ -29,7 +30,7 @@ class Streamer:
 
     def __init__(
         self,
-        trainset: TrainSet,
+        trainset: Union[TrainSetPB, TrainSetModel],
         reader_headers: Dict[str, str],
         base_url: str,
         kbid: str,
@@ -47,12 +48,23 @@ class Streamer:
     def initialize(self, partition_id: str):
         self.stream_session = requests.Session()
         self.stream_session.headers.update(self.reader_headers)
-        self.resp = self.stream_session.post(
-            f"{self.base_url}/v1/kb/{self.kbid}/trainset/{partition_id}",
-            data=self.trainset.SerializeToString(),
-            stream=True,
-            timeout=None,
-        )
+        if isinstance(self.trainset, TrainSetPB):
+            # Legacy version of the endpoint is passing the protobuffer as bytes in the request content
+            self.resp = self.stream_session.post(
+                f"{self.base_url}/v1/kb/{self.kbid}/trainset/{partition_id}",
+                data=self.trainset.SerializeToString(),
+                stream=True,
+                timeout=None,
+            )
+        elif isinstance(self.trainset, TrainSetModel):
+            self.resp = self.stream_session.post(
+                f"{self.base_url}/v1/kb/{self.kbid}/trainset/{partition_id}",
+                json=self.trainset.model_dump(),
+                stream=True,
+                timeout=None,
+            )
+        else:  # pragma: no cover
+            raise ValueError("Invalid trainset type")
         self.resp.raise_for_status()
 
     def finalize(self):

{nucliadb_dataset-6.4.2.post4389.dist-info → nucliadb_dataset-6.5.0.post4404.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb_dataset
-Version: 6.4.2.post4389
+Version: 6.5.0.post4404
 Summary: NucliaDB Train Python client
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: Apache-2.0
@@ -25,9 +25,9 @@ Requires-Dist: aiohttp
 Requires-Dist: argdantic
 Requires-Dist: pydantic-settings>=2.2
 Requires-Dist: pyarrow
-Requires-Dist: nucliadb-protos>=6.4.2.post4389
-Requires-Dist: nucliadb-sdk>=6.4.2.post4389
-Requires-Dist: nucliadb-models>=6.4.2.post4389
+Requires-Dist: nucliadb-protos>=6.5.0.post4404
+Requires-Dist: nucliadb-sdk>=6.5.0.post4404
+Requires-Dist: nucliadb-models>=6.5.0.post4404
 
 # NUCLIADB TRAIN CLIENT
 

{nucliadb_dataset-6.4.2.post4389.dist-info → nucliadb_dataset-6.5.0.post4404.dist-info}/RECORD

@@ -1,18 +1,18 @@
 nucliadb_dataset/__init__.py,sha256=I58PAYrrgLvxmkGGHvzKKUwnaZ2ny44hba6AXEYfKOQ,1054
 nucliadb_dataset/api.py,sha256=RDIW23cy12E5_UlwsiOuhdFrr1OHPA4Mj7hZH0BqGgA,2757
-nucliadb_dataset/dataset.py,sha256=E8ARHwZntdZH9g4zKEVhYbqMYxn3hnuPGkuErL0V-KM,7951
+nucliadb_dataset/dataset.py,sha256=11XoRslzMQQHwJA5MrdAQ30eLP8DoNTv5EmuzL2mln0,8061
 nucliadb_dataset/export.py,sha256=FJjmg1GA0fhzxEZVgMbrqjbcLTi2v9gw3wI0vhKaDWI,2528
 nucliadb_dataset/mapping.py,sha256=Ayg-dDiGc4P-ctRj2ddIlbgidziuLlZGKeAZ08aHBZU,6495
 nucliadb_dataset/nuclia.py,sha256=uXiwJS_GBcN6z9l9hFahD3jLbqhkfpGNWpMR6_8K5k8,3718
 nucliadb_dataset/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb_dataset/run.py,sha256=Ktqv6m0f5oCs54RNZh9b3MJtYZ_JmKP1Zp2Uo3VEyF4,2806
 nucliadb_dataset/settings.py,sha256=9NYeJVgIHbLqCenUFRwj6Iz9S7klOROXQfzaUiBUEl0,2027
-nucliadb_dataset/streamer.py,sha256=-v4FeIRhU_uLhic14vKsbBg8sIqDg9O4gtk_JwSqScg,2381
+nucliadb_dataset/streamer.py,sha256=aBzYWNQQVWM5qA14f8p0YQdshzC4KpOQFt3J0TYe1uk,3060
 nucliadb_dataset/tasks.py,sha256=198o37vDlzS7OdXrHYhtwI8kz2WHWJnxpholh-rtTPQ,6227
 nucliadb_dataset/tests/__init__.py,sha256=zG33bUz1rHFPtvqQPWn4rDwBJt3FJodGuQYD45quiQg,583
 nucliadb_dataset/tests/fixtures.py,sha256=Fx82tfwoFx0YyaiC6tg-ThUDvS8pcmFSE0Pm1KpAS0I,9242
-nucliadb_dataset-6.4.2.post4389.dist-info/METADATA,sha256=wB4H4rkEt3JZQ4Ec3yxLVAHgTGUOQOHGqLG9LyKDEh4,1263
-nucliadb_dataset-6.4.2.post4389.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb_dataset-6.4.2.post4389.dist-info/entry_points.txt,sha256=ORrWnn6AUFfHGY1fWPRPTz99KV_pXXwttZAopyT8qvQ,60
-nucliadb_dataset-6.4.2.post4389.dist-info/top_level.txt,sha256=aJtDe54tz6060E0uyk1rdTRAU4FPWo5if1fYFQGvdqU,17
-nucliadb_dataset-6.4.2.post4389.dist-info/RECORD,,
+nucliadb_dataset-6.5.0.post4404.dist-info/METADATA,sha256=wr3_UeG9eunFrhgNc3GRmBFZ9x_EoFhYwm48T5TmaSg,1263
+nucliadb_dataset-6.5.0.post4404.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb_dataset-6.5.0.post4404.dist-info/entry_points.txt,sha256=ORrWnn6AUFfHGY1fWPRPTz99KV_pXXwttZAopyT8qvQ,60
+nucliadb_dataset-6.5.0.post4404.dist-info/top_level.txt,sha256=aJtDe54tz6060E0uyk1rdTRAU4FPWo5if1fYFQGvdqU,17
+nucliadb_dataset-6.5.0.post4404.dist-info/RECORD,,