nucliadb-dataset 6.10.0.post5655__py3-none-any.whl → 6.10.0.post5658__py3-none-any.whl

nucliadb_dataset/__init__.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 
 from enum import Enum
-from typing import Dict
 
 from nucliadb_dataset.dataset import NucliaDBDataset, Task, download_all_partitions
 from nucliadb_dataset.nuclia import NucliaDriver
 
-NUCLIA_GLOBAL: Dict[str, NucliaDriver] = {}
+NUCLIA_GLOBAL: dict[str, NucliaDriver] = {}
 
 CLIENT_ID = "CLIENT"
 
@@ -29,10 +28,10 @@ class ExportType(str, Enum):
 
 
 __all__ = (
+    "CLIENT_ID",
+    "NUCLIA_GLOBAL",
+    "ExportType",
     "NucliaDBDataset",
     "Task",
     "download_all_partitions",
-    "NUCLIA_GLOBAL",
-    "CLIENT_ID",
-    "ExportType",
 )

nucliadb_dataset/api.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Iterator, List
+from collections.abc import Iterator
 
 from nucliadb_dataset import CLIENT_ID, NUCLIA_GLOBAL
 from nucliadb_dataset.nuclia import NucliaDriver
@@ -36,26 +36,22 @@ def get_nuclia_client() -> NucliaDriver:
 
 def iterate_sentences(kbid: str, labels: bool, entities: bool, text: bool) -> Iterator[TrainSentence]:
     client = get_nuclia_client()
-    for sentence in client.iterate_sentences(kbid, labels, entities, text):
-        yield sentence
+    yield from client.iterate_sentences(kbid, labels, entities, text)
 
 
 def iterate_paragraphs(kbid: str, labels: bool, entities: bool, text: bool) -> Iterator[TrainParagraph]:
     client = get_nuclia_client()
-    for sentence in client.iterate_paragraphs(kbid, labels, entities, text):
-        yield sentence
+    yield from client.iterate_paragraphs(kbid, labels, entities, text)
 
 
 def iterate_fields(kbid: str, labels: bool, entities: bool, text: bool) -> Iterator[TrainField]:
     client = get_nuclia_client()
-    for sentence in client.iterate_fields(kbid, labels, entities, text):
-        yield sentence
+    yield from client.iterate_fields(kbid, labels, entities, text)
 
 
 def iterate_resources(kbid: str, labels: bool, entities: bool, text: bool) -> Iterator[TrainResource]:
     client = get_nuclia_client()
-    for sentence in client.iterate_resources(kbid, labels, entities, text):
-        yield sentence
+    yield from client.iterate_resources(kbid, labels, entities, text)
 
 
 def get_labels(kbid: str) -> GetLabelsResponse:
@@ -76,7 +72,7 @@ def get_info(kbid: str) -> TrainInfo:
     return info
 
 
-def get_ontology_count(kbid: str, paragraph_labelsets: List[str], resource_labelset: List[str]):
+def get_ontology_count(kbid: str, paragraph_labelsets: list[str], resource_labelset: list[str]):
     client = get_nuclia_client()
     labels = client.get_ontology_count(kbid, paragraph_labelsets, resource_labelset)
     return labels

nucliadb_dataset/dataset.py

@@ -14,8 +14,9 @@
 
 import logging
 import os
+from collections.abc import Callable, Iterator
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
+from typing import Any
 
 import pyarrow as pa  # type: ignore
 
@@ -42,12 +43,12 @@ CHUNK_SIZE = 5 * 1024 * 1024
 @dataclass
 class LabelSetCount:
     count: int
-    labels: Dict[str, int] = field(default_factory=dict)
+    labels: dict[str, int] = field(default_factory=dict)
 
 
-class NucliaDataset(object):
-    labels: Optional[KnowledgeBoxLabels]
-    entities: Optional[KnowledgeBoxEntities]
+class NucliaDataset:
+    labels: KnowledgeBoxLabels | None
+    entities: KnowledgeBoxEntities | None
 
     def __new__(cls, *args, **kwargs):
         if cls is NucliaDataset:
@@ -56,18 +57,18 @@ class NucliaDataset(object):
 
     def __init__(
         self,
-        base_path: Optional[str] = None,
+        base_path: str | None = None,
     ):
         if base_path is None:
             base_path = os.getcwd()
         self.base_path = base_path
-        self.mappings: List[Callable] = []
+        self.mappings: list[Callable] = []
 
         self.labels = None
         self.entities = None
         self.folder = None
 
-    def iter_all_partitions(self, force=False) -> Iterator[Tuple[str, str]]:
+    def iter_all_partitions(self, force=False) -> Iterator[tuple[str, str]]:
         partitions = self.get_partitions()
         for index, partition in enumerate(partitions):
             logger.info(f"Reading partition {partition} {index}/{len(partitions)}")
@@ -75,7 +76,7 @@ class NucliaDataset(object):
             logger.info(f"Done reading partition {partition}")
             yield partition, filename
 
-    def read_all_partitions(self, force=False, path: Optional[str] = None) -> List[str]:
+    def read_all_partitions(self, force=False, path: str | None = None) -> list[str]:
         partitions = self.get_partitions()
         result = []
         for index, partition in enumerate(partitions):
@@ -91,9 +92,9 @@ class NucliaDataset(object):
     def read_partition(
         self,
         partition_id: str,
-        filename: Optional[str] = None,
+        filename: str | None = None,
         force: bool = False,
-        path: Optional[str] = None,
+        path: str | None = None,
     ):
         raise NotImplementedError()
 
@@ -103,12 +104,12 @@ class NucliaDBDataset(NucliaDataset):
         self,
         sdk: NucliaDB,
         kbid: str,
-        task: Optional[Task] = None,
-        labels: Optional[List[str]] = None,
-        trainset: Optional[Union[TrainSetPB, TrainSetModel]] = None,
-        base_path: Optional[str] = None,
-        search_sdk: Optional[NucliaDB] = None,
-        reader_sdk: Optional[NucliaDB] = None,
+        task: Task | None = None,
+        labels: list[str] | None = None,
+        trainset: TrainSetPB | TrainSetModel | None = None,
+        base_path: str | None = None,
+        search_sdk: NucliaDB | None = None,
+        reader_sdk: NucliaDB | None = None,
     ):
         super().__init__(base_path)
 
@@ -165,13 +166,13 @@ class NucliaDBDataset(NucliaDataset):
         streamer.initialize(partition_id)
         return streamer
 
-    def _set_mappings(self, funcs: List[Callable[[Any, Any], Tuple[Any, Any]]]):
+    def _set_mappings(self, funcs: list[Callable[[Any, Any], tuple[Any, Any]]]):
         self.mappings = funcs
 
     def _set_schema(self, schema: pa.Schema):
         self.schema = schema
 
-    def get_partitions(self) -> List[str]:
+    def get_partitions(self) -> list[str]:
         """
         Get expected number of partitions from a live NucliaDB
         """
@@ -183,9 +184,9 @@ class NucliaDBDataset(NucliaDataset):
     def read_partition(
         self,
         partition_id: str,
-        filename: Optional[str] = None,
+        filename: str | None = None,
         force: bool = False,
-        path: Optional[str] = None,
+        path: str | None = None,
     ):
         """
         Export an arrow partition from a live NucliaDB and store it locally
@@ -220,12 +221,12 @@ class NucliaDBDataset(NucliaDataset):
 
 def download_all_partitions(
     task: str,
-    slug: Optional[str] = None,
-    kbid: Optional[str] = None,
-    nucliadb_base_url: Optional[str] = "http://localhost:8080",
-    path: Optional[str] = None,
-    sdk: Optional[NucliaDB] = None,
-    labels: Optional[List[str]] = None,
+    slug: str | None = None,
+    kbid: str | None = None,
+    nucliadb_base_url: str | None = "http://localhost:8080",
+    path: str | None = None,
+    sdk: NucliaDB | None = None,
+    labels: list[str] | None = None,
 ):
     if sdk is None:
         sdk = NucliaDB(region="on-prem", url=nucliadb_base_url)

nucliadb_dataset/nuclia.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Iterator, List
+from collections.abc import Iterator
 
 import grpc
 
@@ -54,8 +54,7 @@ class NucliaDriver:
         request.metadata.labels = labels
         request.metadata.entities = entities
         request.metadata.text = text
-        for sentence in self.stub.GetSentences(request):
-            yield sentence
+        yield from self.stub.GetSentences(request)
 
     def iterate_paragraphs(
         self, kbid: str, labels: bool, entities: bool, text: bool
@@ -65,8 +64,7 @@ class NucliaDriver:
         request.metadata.labels = labels
         request.metadata.entities = entities
         request.metadata.text = text
-        for paragraph in self.stub.GetParagraphs(request):
-            yield paragraph
+        yield from self.stub.GetParagraphs(request)
 
     def iterate_resources(
         self, kbid: str, labels: bool, entities: bool, text: bool
@@ -76,8 +74,7 @@ class NucliaDriver:
         request.metadata.labels = labels
         request.metadata.entities = entities
         request.metadata.text = text
-        for resource in self.stub.GetResources(request):
-            yield resource
+        yield from self.stub.GetResources(request)
 
     def iterate_fields(
         self, kbid: str, labels: bool, entities: bool, text: bool
@@ -87,8 +84,7 @@ class NucliaDriver:
         request.metadata.labels = labels
         request.metadata.entities = entities
         request.metadata.text = text
-        for field in self.stub.GetFields(request):
-            yield field
+        yield from self.stub.GetFields(request)
 
     def get_labels(self, kbid: str) -> GetLabelsResponse:
         request = GetLabelsRequest()
@@ -106,7 +102,7 @@ class NucliaDriver:
         return self.stub.GetInfo(request)
 
     def get_ontology_count(
-        self, kbid: str, paragraph_labelsets: List[str], resource_labelsets: List[str]
+        self, kbid: str, paragraph_labelsets: list[str], resource_labelsets: list[str]
     ) -> LabelsetsCount:
         request = GetLabelsetsCountRequest()
         request.kb.uuid = kbid

nucliadb_dataset/settings.py

@@ -14,7 +14,6 @@
 #
 
 from pathlib import Path
-from typing import Optional
 
 import pydantic
 from pydantic_settings import BaseSettings
@@ -35,20 +34,20 @@ class RunningSettings(BaseSettings):
     download_path: str = pydantic.Field(f"{Path.home()}/.nuclia/download", description="Download path")
     url: str = pydantic.Field(description="KnowledgeBox URL")
     type: Task = pydantic.Field(description="Dataset Type")
-    labelset: Optional[str] = pydantic.Field(
+    labelset: str | None = pydantic.Field(
         None, description="For classification which labelset or families"
     )
 
     datasets_url: str = pydantic.Field(
         "https://europe-1.nuclia.cloud",
-        description="Base url for the Nuclia datasets component (excluding /api/v1)™",  # noqa
+        description="Base url for the Nuclia datasets component (excluding /api/v1)™",
     )
 
-    apikey: Optional[str] = pydantic.Field(None, description="API key to upload to Nuclia Datasets™")
+    apikey: str | None = pydantic.Field(None, description="API key to upload to Nuclia Datasets™")
 
     environment: str = pydantic.Field("on-prem", description="region or on-prem")
 
-    service_token: Optional[str] = pydantic.Field(
+    service_token: str | None = pydantic.Field(
         None, description="Service account key to access Nuclia Cloud"
     )
 

nucliadb_dataset/streamer.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import logging
-from typing import Dict, Optional, Union
 
 import requests
 
@@ -26,12 +25,12 @@ SIZE_BYTES = 4
 
 
 class Streamer:
-    resp: Optional[requests.Response]
+    resp: requests.Response | None
 
     def __init__(
         self,
-        trainset: Union[TrainSetPB, TrainSetModel],
-        reader_headers: Dict[str, str],
+        trainset: TrainSetPB | TrainSetModel,
+        reader_headers: dict[str, str],
         base_url: str,
         kbid: str,
     ):
@@ -75,16 +74,16 @@ class Streamer:
     def __iter__(self):
         return self
 
-    def read(self) -> Optional[bytes]:
+    def read(self) -> bytes | None:
         assert self.resp is not None, "Streamer not initialized"
         header = self.resp.raw.read(4, decode_content=True)
         if header == b"":
             return None
-        payload_size = int.from_bytes(header, byteorder="big", signed=False)  # noqa
+        payload_size = int.from_bytes(header, byteorder="big", signed=False)
         data = self.resp.raw.read(payload_size)
         return data
 
-    def __next__(self) -> Optional[bytes]:
+    def __next__(self) -> bytes | None:
         payload = self.read()
         if payload in [None, b""]:
             logger.info("Streamer finished reading")

nucliadb_dataset/tasks.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Callable
 from dataclasses import dataclass
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable, Dict, List
+from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa  # type: ignore
 
@@ -64,10 +65,10 @@ class TaskDefinition:
     schema: pa.schema
     proto: Any
     labels: bool
-    mapping: List[Callable]
+    mapping: list[Callable]
 
 
-TASK_DEFINITIONS: Dict[Task, TaskDefinition] = {
+TASK_DEFINITIONS: dict[Task, TaskDefinition] = {
     Task.PARAGRAPH_CLASSIFICATION: TaskDefinition(
         schema=pa.schema(
             [
@@ -190,4 +191,4 @@ TASK_DEFINITIONS: Dict[Task, TaskDefinition] = {
     ),
 }
 
-TASK_DEFINITIONS_REVERSE = {task.proto: task for task in TASK_DEFINITIONS.values()}  # noqa
+TASK_DEFINITIONS_REVERSE = {task.proto: task for task in TASK_DEFINITIONS.values()}

nucliadb_dataset/tests/fixtures.py

@@ -14,7 +14,8 @@
 
 import re
 import tempfile
-from typing import TYPE_CHECKING, AsyncIterator, Iterator, Optional
+from collections.abc import AsyncIterator, Iterator
+from typing import TYPE_CHECKING
 
 import docker  # type: ignore
 import grpc
@@ -32,7 +33,7 @@ from nucliadb_protos.writer_pb2_grpc import WriterStub
 from nucliadb_sdk.v2.sdk import NucliaDB
 
 DOCKER_ENV_GROUPS = re.search(r"//([^:]+)", docker.from_env().api.base_url)
-DOCKER_HOST: Optional[str] = DOCKER_ENV_GROUPS.group(1) if DOCKER_ENV_GROUPS else None
+DOCKER_HOST: str | None = DOCKER_ENV_GROUPS.group(1) if DOCKER_ENV_GROUPS else None
 
 if TYPE_CHECKING:
     from nucliadb_protos.writer_pb2_grpc import WriterAsyncStub

{nucliadb_dataset-6.10.0.post5655.dist-info → nucliadb_dataset-6.10.0.post5658.dist-info}/METADATA

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

nucliadb_dataset-6.10.0.post5658.dist-info/RECORD

@@ -0,0 +1,18 @@
+nucliadb_dataset/__init__.py,sha256=1lvjYSji93zdS2zZGnanh9TeunJcaInv-CBEoEfdAL0,1030
+nucliadb_dataset/api.py,sha256=i8OlwF1ly6lkrEMb9Ffoc0GSXB-7zGaOtXMN51wDyO8,2644
+nucliadb_dataset/dataset.py,sha256=VeJH71WWNJPFO6UIKRP8iSz5O6FDXwOUaQdS52VuTNQ,7978
+nucliadb_dataset/export.py,sha256=FJjmg1GA0fhzxEZVgMbrqjbcLTi2v9gw3wI0vhKaDWI,2528
+nucliadb_dataset/mapping.py,sha256=Ayg-dDiGc4P-ctRj2ddIlbgidziuLlZGKeAZ08aHBZU,6495
+nucliadb_dataset/nuclia.py,sha256=eBsuDs_HHZTpbyl35H7nqqu3pgEFHzDo8VygRFjTWDg,3593
+nucliadb_dataset/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nucliadb_dataset/run.py,sha256=Ktqv6m0f5oCs54RNZh9b3MJtYZ_JmKP1Zp2Uo3VEyF4,2806
+nucliadb_dataset/settings.py,sha256=xRDTuqj4SVyTmliA6RwMRmbvI4aLXTx-CrfUDqMJvbc,1982
+nucliadb_dataset/streamer.py,sha256=EKoasniFuVHSQIPVxZ6i4wQVKsXxDNLtPjpX_J5CLDk,2996
+nucliadb_dataset/tasks.py,sha256=e1pCPzY9T2jQP2WgTNHGJnuBKg0JWg0vonuxQWAnJLs,6234
+nucliadb_dataset/tests/__init__.py,sha256=zG33bUz1rHFPtvqQPWn4rDwBJt3FJodGuQYD45quiQg,583
+nucliadb_dataset/tests/fixtures.py,sha256=KPMO0mYzX3HR41nzkJ8D0TlF7bUekQGBi8uSaMtAmLU,7524
+nucliadb_dataset-6.10.0.post5658.dist-info/METADATA,sha256=bi5YmKleoczubQfykl0Mgnh9PmlABjmoaUodIji-D5Y,1218
+nucliadb_dataset-6.10.0.post5658.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb_dataset-6.10.0.post5658.dist-info/entry_points.txt,sha256=ORrWnn6AUFfHGY1fWPRPTz99KV_pXXwttZAopyT8qvQ,60
+nucliadb_dataset-6.10.0.post5658.dist-info/top_level.txt,sha256=aJtDe54tz6060E0uyk1rdTRAU4FPWo5if1fYFQGvdqU,17
+nucliadb_dataset-6.10.0.post5658.dist-info/RECORD,,

nucliadb_dataset-6.10.0.post5655.dist-info/RECORD

@@ -1,18 +0,0 @@
-nucliadb_dataset/__init__.py,sha256=I58PAYrrgLvxmkGGHvzKKUwnaZ2ny44hba6AXEYfKOQ,1054
-nucliadb_dataset/api.py,sha256=RDIW23cy12E5_UlwsiOuhdFrr1OHPA4Mj7hZH0BqGgA,2757
-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=aBzYWNQQVWM5qA14f8p0YQdshzC4KpOQFt3J0TYe1uk,3060
-nucliadb_dataset/tasks.py,sha256=198o37vDlzS7OdXrHYhtwI8kz2WHWJnxpholh-rtTPQ,6227
-nucliadb_dataset/tests/__init__.py,sha256=zG33bUz1rHFPtvqQPWn4rDwBJt3FJodGuQYD45quiQg,583
-nucliadb_dataset/tests/fixtures.py,sha256=5Ac20TKYmGe61d74yRT6hj0ENdLr3jT2Gbj2fiSX9LM,7510
-nucliadb_dataset-6.10.0.post5655.dist-info/METADATA,sha256=OntzvxWPCKnnI0_7fnqZVCE0sdOewphARAbwOu_MPVk,1218
-nucliadb_dataset-6.10.0.post5655.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb_dataset-6.10.0.post5655.dist-info/entry_points.txt,sha256=ORrWnn6AUFfHGY1fWPRPTz99KV_pXXwttZAopyT8qvQ,60
-nucliadb_dataset-6.10.0.post5655.dist-info/top_level.txt,sha256=aJtDe54tz6060E0uyk1rdTRAU4FPWo5if1fYFQGvdqU,17
-nucliadb_dataset-6.10.0.post5655.dist-info/RECORD,,