nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/train/generators/question_answer_streaming.py

@@ -20,6 +20,14 @@
 
 from typing import AsyncGenerator
 
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.ids import FIELD_TYPE_PB_TO_STR, FIELD_TYPE_STR_TO_PB
+from nucliadb.train import logger
+from nucliadb.train.generators.utils import (
+    batchify,
+    get_paragraph,
+    get_resource_from_cache_or_db,
+)
 from nucliadb_protos.dataset_pb2 import (
     QuestionAnswerStreamingBatch,
     QuestionAnswerStreamItem,
@@ -32,15 +40,6 @@ from nucliadb_protos.resources_pb2 import (
     QuestionAnswerAnnotation,
 )
 
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.ingest.orm.resource import FIELD_TYPE_TO_ID, KB_REVERSE
-from nucliadb.train import logger
-from nucliadb.train.generators.utils import (
-    batchify,
-    get_paragraph,
-    get_resource_from_cache_or_db,
-)
-
 
 def question_answer_batch_generator(
     kbid: str,
@@ -48,12 +47,8 @@ def question_answer_batch_generator(
     node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[QuestionAnswerStreamingBatch, None]:
-    generator = generate_question_answer_streaming_payloads(
-        kbid, trainset, node, shard_replica_id
-    )
-    batch_generator = batchify(
-        generator, trainset.batch_size, QuestionAnswerStreamingBatch
-    )
+    generator = generate_question_answer_streaming_payloads(kbid, trainset, node, shard_replica_id)
+    batch_generator = batchify(generator, trainset.batch_size, QuestionAnswerStreamingBatch)
     return batch_generator
 
 
@@ -90,14 +85,18 @@ async def generate_question_answer_streaming_payloads(
                         item.cancelled_by_user = qa_annotation_pb.cancelled_by_user
                         yield item
 
-        field_type_int = KB_REVERSE[field_type]
+        field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
         field_obj = await orm_resource.get_field(field, field_type_int, load=False)
 
         question_answers_pb = await field_obj.get_question_answers()
         if question_answers_pb is not None:
-            for question_answer_pb in question_answers_pb.question_answer:
+            for question_answer_pb in question_answers_pb.question_answers.question_answer:
                 async for item in iter_stream_items(kbid, question_answer_pb):
                     yield item
+            for question_answer_pb in question_answers_pb.split_question_answers.values():
+                for split_question_answer_pb in question_answer_pb.question_answer:
+                    async for item in iter_stream_items(kbid, split_question_answer_pb):
+                        yield item
 
 
 async def iter_stream_items(
@@ -109,7 +108,7 @@ async def iter_stream_items(
     for paragraph_id in question_pb.ids_paragraphs:
         try:
             text = await get_paragraph(kbid, paragraph_id)
-        except Exception as exc:  # pragma: nocover
+        except Exception as exc:  # pragma: no cover
             logger.warning(
                 "Question paragraph couldn't be fetched while streaming Q&A",
                 extra={"kbid": kbid, "paragraph_id": paragraph_id},
@@ -128,7 +127,7 @@ async def iter_stream_items(
         for paragraph_id in answer_pb.ids_paragraphs:
             try:
                 text = await get_paragraph(kbid, paragraph_id)
-            except Exception as exc:  # pragma: nocover
+            except Exception as exc:  # pragma: no cover
                 logger.warning(
                     "Answer paragraph couldn't be fetched while streaming Q&A",
                     extra={"kbid": kbid, "paragraph_id": paragraph_id},
@@ -141,4 +140,4 @@ async def iter_stream_items(
 
 
 def is_same_field(field: FieldID, field_id: str, field_type: str) -> bool:
-    return field.field == field_id and FIELD_TYPE_TO_ID[field.field_type] == field_type
+    return field.field == field_id and FIELD_TYPE_PB_TO_STR[field.field_type] == field_type

nucliadb/train/generators/sentence_classifier.py

@@ -21,6 +21,11 @@
 from typing import AsyncGenerator
 
 from fastapi import HTTPException
+
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.train import logger
+from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
     Label,
     MultipleTextSameLabels,
@@ -29,11 +34,6 @@ from nucliadb_protos.dataset_pb2 import (
 )
 from nucliadb_protos.nodereader_pb2 import StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.ingest.orm.resource import KB_REVERSE
-from nucliadb.train import logger
-from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
-
 
 def sentence_classification_batch_generator(
     kbid: str,
@@ -47,12 +47,8 @@ def sentence_classification_batch_generator(
             detail="Sentence Classification should be at least of 1 labelset",
         )
 
-    generator = generate_sentence_classification_payloads(
-        kbid, trainset, node, shard_replica_id
-    )
-    batch_generator = batchify(
-        generator, trainset.batch_size, SentenceClassificationBatch
-    )
+    generator = generate_sentence_classification_payloads(kbid, trainset, node, shard_replica_id)
+    batch_generator = batchify(generator, trainset.batch_size, SentenceClassificationBatch)
     return batch_generator
 
 
@@ -107,14 +103,12 @@ async def get_sentences(kbid: str, result: str) -> list[str]:
         logger.error(f"{rid} does not exist on DB")
         return []
 
-    field_type_int = KB_REVERSE[field_type]
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
     extracted_text = await field_obj.get_extracted_text()
     field_metadata = await field_obj.get_field_metadata()
     if extracted_text is None:
-        logger.warning(
-            f"{rid} {field} {field_type_int} extracted_text does not exist on DB"
-        )
+        logger.warning(f"{rid} {field} {field_type_int} extracted_text does not exist on DB")
         return []
 
     splitted_texts = []

nucliadb/train/generators/token_classifier.py

@@ -21,6 +21,10 @@
 from collections import OrderedDict
 from typing import AsyncGenerator, cast
 
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.train import logger
+from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
     TokenClassificationBatch,
     TokensClassification,
@@ -28,11 +32,6 @@ from nucliadb_protos.dataset_pb2 import (
 )
 from nucliadb_protos.nodereader_pb2 import StreamFilter, StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.ingest.orm.resource import KB_REVERSE
-from nucliadb.train import logger
-from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
-
 NERS_DICT = dict[str, dict[str, list[tuple[int, int]]]]
 POSITION_DICT = OrderedDict[tuple[int, int], tuple[str, str]]
 MAIN = "__main__"
@@ -44,9 +43,7 @@ def token_classification_batch_generator(
     node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[TokenClassificationBatch, None]:
-    generator = generate_token_classification_payloads(
-        kbid, trainset, node, shard_replica_id
-    )
+    generator = generate_token_classification_payloads(kbid, trainset, node, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, TokenClassificationBatch)
     return batch_generator
 
@@ -97,13 +94,11 @@ async def get_field_text(
         logger.error(f"{rid} does not exist on DB")
         return {}, {}, {}
 
-    field_type_int = KB_REVERSE[field_type]
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
     extracted_text = await field_obj.get_extracted_text()
     if extracted_text is None:
-        logger.warning(
-            f"{rid} {field} {field_type_int} extracted_text does not exist on DB"
-        )
+        logger.warning(f"{rid} {field} {field_type_int} extracted_text does not exist on DB")
         return {}, {}, {}
 
     split_text: dict[str, str] = extracted_text.split_text
@@ -138,16 +133,42 @@ async def get_field_text(
                                 split = MAIN
                             else:
                                 split = token.split
-                            split_ners[split].setdefault(token.klass, {}).setdefault(
-                                token.token, []
-                            )
-                            split_ners[split][token.klass][token.token].append(
-                                (token.start, token.end)
-                            )
+                            split_ners[split].setdefault(token.klass, {}).setdefault(token.token, [])
+                            split_ners[split][token.klass][token.token].append((token.start, token.end))
 
     field_metadata = await field_obj.get_field_metadata()
     # Check computed definition of entities
     if field_metadata is not None:
+        # Data Augmentation + Processor entities
+        for data_augmentation_task_id, entities in field_metadata.metadata.entities.items():
+            for entity in entities.entities:
+                entity_text = entity.text
+                entity_label = entity.label
+                entity_positions = entity.positions
+                if entity_label in valid_entity_groups:
+                    split_ners[MAIN].setdefault(entity_label, {}).setdefault(entity_text, [])
+                    for position in entity_positions:
+                        split_ners[MAIN][entity_label][entity_text].append(
+                            (position.start, position.end)
+                        )
+
+        for split, split_metadata in field_metadata.split_metadata.items():
+            for data_augmentation_task_id, entities in split_metadata.entities.items():
+                for entity in entities.entities:
+                    entity_text = entity.text
+                    entity_label = entity.label
+                    entity_positions = entity.positions
+                    if entity_label in valid_entity_groups:
+                        split_ners.setdefault(split, {}).setdefault(entity_label, {}).setdefault(
+                            entity_text, []
+                        )
+                        for position in entity_positions:
+                            split_ners[split][entity_label][entity_text].append(
+                                (position.start, position.end)
+                            )
+
+        # Legacy processor entities
+        # TODO: Remove once processor doesn't use this anymore and remove the positions and ner fields from the message
         for entity_key, positions in field_metadata.metadata.positions.items():
             entities = entity_key.split("/")
             entity_group = entities[0]
@@ -156,9 +177,7 @@ async def get_field_text(
             if entity_group in valid_entity_groups:
                 split_ners[MAIN].setdefault(entity_group, {}).setdefault(entity, [])
                 for position in positions.position:
-                    split_ners[MAIN][entity_group][entity].append(
-                        (position.start, position.end)
-                    )
+                    split_ners[MAIN][entity_group][entity].append((position.start, position.end))
 
         for split, split_metadata in field_metadata.split_metadata.items():
             for entity_key, positions in split_metadata.positions.items():
@@ -166,24 +185,16 @@ async def get_field_text(
                 entity_group = entities[0]
                 entity = "/".join(entities[1:])
                 if entity_group in valid_entity_groups:
-                    split_ners.setdefault(split, {}).setdefault(
-                        entity_group, {}
-                    ).setdefault(entity, [])
+                    split_ners.setdefault(split, {}).setdefault(entity_group, {}).setdefault(entity, [])
                     for position in positions.position:
-                        split_ners[split][entity_group][entity].append(
-                            (position.start, position.end)
-                        )
+                        split_ners[split][entity_group][entity].append((position.start, position.end))
 
     for split, invalid_tokens in invalid_tokens_split.items():
         for token.klass, token.token, token.start, token.end in invalid_tokens:
             if token.klass in split_ners.get(split, {}):
                 if token.token in split_ners.get(split, {}).get(token.klass, {}):
-                    if (token.start, token.end) in split_ners[split][token.klass][
-                        token.token
-                    ]:
-                        split_ners[split][token.klass][token.token].remove(
-                            (token.start, token.end)
-                        )
+                    if (token.start, token.end) in split_ners[split][token.klass][token.token]:
+                        split_ners[split][token.klass][token.token].remove((token.start, token.end))
                         if len(split_ners[split][token.klass][token.token]) == 0:
                             del split_ners[split][token.klass][token.token]
                         if len(split_ners[split][token.klass]) == 0:
@@ -197,9 +208,7 @@ async def get_field_text(
                 for position in positions:
                     split_positions[position] = (entity_group, entity)
 
-        ordered_positions[split] = OrderedDict(
-            sorted(split_positions.items(), key=lambda x: x[0])
-        )
+        ordered_positions[split] = OrderedDict(sorted(split_positions.items(), key=lambda x: x[0]))
 
     split_paragraphs: dict[str, list[tuple[int, int]]] = {}
     if field_metadata is not None:

nucliadb/train/generators/utils.py

@@ -19,19 +19,17 @@
 #
 
 from contextvars import ContextVar
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncGenerator, AsyncIterator, Optional, Type
 
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
-from nucliadb.ingest.orm.resource import KB_REVERSE
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
-from nucliadb.middleware.transaction import get_read_only_transaction
 from nucliadb.train import SERVICE_NAME, logger
-from nucliadb.train.types import TrainBatchType
+from nucliadb.train.types import T
 from nucliadb_utils.utilities import get_storage
 
-rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar(
-    "rcache", default=None
-)
+rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar("rcache", default=None)
 
 
 def get_resource_cache(clear: bool = False) -> dict[str, ResourceORM]:
@@ -46,12 +44,12 @@ async def get_resource_from_cache_or_db(kbid: str, uuid: str) -> Optional[Resour
     resouce_cache = get_resource_cache()
     orm_resource: Optional[ResourceORM] = None
     if uuid not in resouce_cache:
-        transaction = await get_read_only_transaction()
         storage = await get_storage(service_name=SERVICE_NAME)
-        kb = KnowledgeBoxORM(transaction, storage, kbid)
-        orm_resource = await kb.get(uuid)
-        if orm_resource is not None:
-            resouce_cache[uuid] = orm_resource
+        async with get_driver().transaction(read_only=True) as transaction:
+            kb = KnowledgeBoxORM(transaction, storage, kbid)
+            orm_resource = await kb.get(uuid)
+            if orm_resource is not None:
+                resouce_cache[uuid] = orm_resource
     else:
         orm_resource = resouce_cache.get(uuid)
     return orm_resource
@@ -75,13 +73,11 @@ async def get_paragraph(kbid: str, paragraph_id: str) -> str:
         logger.error(f"{rid} does not exist on DB")
         return ""
 
-    field_type_int = KB_REVERSE[field_type]
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
     extracted_text = await field_obj.get_extracted_text()
     if extracted_text is None:
-        logger.warning(
-            f"{rid} {field} {field_type_int} extracted_text does not exist on DB"
-        )
+        logger.warning(f"{rid} {field} {field_type_int} extracted_text does not exist on DB")
         return ""
 
     if split is not None:
@@ -94,8 +90,8 @@ async def get_paragraph(kbid: str, paragraph_id: str) -> str:
 
 
 async def batchify(
-    producer: AsyncIterator[Any], size: int, batch_klass: TrainBatchType
-):
+    producer: AsyncIterator[Any], size: int, batch_klass: Type[T]
+) -> AsyncGenerator[T, None]:
     # NOTE: we are supposing all protobuffers have a data field
     batch = []
     async for item in producer:

nucliadb/train/lifecycle.py

@@ -18,6 +18,10 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
 from nucliadb.common.cluster.discovery.utils import (
     setup_cluster_discovery,
     teardown_cluster_discovery,
@@ -33,16 +37,16 @@ from nucliadb_telemetry.utils import clean_telemetry, setup_telemetry
 from nucliadb_utils.utilities import start_audit_utility, stop_audit_utility
 
 
-async def initialize() -> None:
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     await setup_telemetry(SERVICE_NAME)
-
     await setup_cluster_discovery()
     await start_shard_manager()
     await start_train_grpc(SERVICE_NAME)
     await start_audit_utility(SERVICE_NAME)
 
+    yield
 
-async def finalize() -> None:
     await stop_audit_utility()
     await stop_train_grpc()
     await stop_shard_manager()

nucliadb/train/nodes.py

@@ -19,6 +19,15 @@
 #
 from typing import AsyncIterator, Optional
 
+from nucliadb.common import datamanagers
+from nucliadb.common.cluster import manager
+from nucliadb.common.cluster.base import AbstractIndexNode
+
+# XXX: this keys shouldn't be exposed outside datamanagers
+from nucliadb.common.datamanagers.resources import KB_RESOURCE_SLUG_BASE
+from nucliadb.common.maindb.driver import Driver, Transaction
+from nucliadb.ingest.orm.entities import EntitiesManager
+from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb_protos.train_pb2 import (
     GetFieldsRequest,
     GetParagraphsRequest,
@@ -30,15 +39,9 @@ from nucliadb_protos.train_pb2 import (
     TrainSentence,
 )
 from nucliadb_protos.writer_pb2 import ShardObject
-
-from nucliadb.common import datamanagers
-from nucliadb.common.cluster import manager
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.common.maindb.driver import Driver, Transaction
-from nucliadb.ingest.orm.entities import EntitiesManager
-from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
-from nucliadb.ingest.orm.resource import KB_RESOURCE_SLUG_BASE
+from nucliadb_utils import const
 from nucliadb_utils.storages.storage import Storage
+from nucliadb_utils.utilities import has_feature
 
 
 class TrainShardManager(manager.KBShardManager):
@@ -50,13 +53,13 @@ class TrainShardManager(manager.KBShardManager):
     async def get_reader(self, kbid: str, shard: str) -> tuple[AbstractIndexNode, str]:
         shards = await self.get_shards_by_kbid_inner(kbid)
         try:
-            shard_object: ShardObject = next(
-                filter(lambda x: x.shard == shard, shards.shards)
-            )
+            shard_object: ShardObject = next(filter(lambda x: x.shard == shard, shards.shards))
         except StopIteration:
             raise KeyError("Shard not found")
 
-        node_obj, shard_id = manager.choose_node(shard_object)
+        node_obj, shard_id = manager.choose_node(
+            shard_object, use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": kbid})
+        )
         return node_obj, shard_id
 
     async def get_kb_obj(self, txn: Transaction, kbid: str) -> Optional[KnowledgeBox]:
@@ -69,9 +72,7 @@ class TrainShardManager(manager.KBShardManager):
         kbobj = KnowledgeBox(txn, self.storage, kbid)
         return kbobj
 
-    async def get_kb_entities_manager(
-        self, txn: Transaction, kbid: str
-    ) -> Optional[EntitiesManager]:
+    async def get_kb_entities_manager(self, txn: Transaction, kbid: str) -> Optional[EntitiesManager]:
         kbobj = await self.get_kb_obj(txn, kbid)
         if kbobj is None:
             return None
@@ -79,9 +80,7 @@ class TrainShardManager(manager.KBShardManager):
         manager = EntitiesManager(kbobj, txn)
         return manager
 
-    async def kb_sentences(
-        self, request: GetSentencesRequest
-    ) -> AsyncIterator[TrainSentence]:
+    async def kb_sentences(self, request: GetSentencesRequest) -> AsyncIterator[TrainSentence]:
         async with self.driver.transaction() as txn:
             kb = KnowledgeBox(txn, self.storage, request.kb.uuid)
             if request.uuid != "":
@@ -95,24 +94,18 @@ class TrainShardManager(manager.KBShardManager):
                     async for sentence in resource.iterate_sentences(request.metadata):
                         yield sentence
 
-    async def kb_paragraphs(
-        self, request: GetParagraphsRequest
-    ) -> AsyncIterator[TrainParagraph]:
+    async def kb_paragraphs(self, request: GetParagraphsRequest) -> AsyncIterator[TrainParagraph]:
         async with self.driver.transaction() as txn:
             kb = KnowledgeBox(txn, self.storage, request.kb.uuid)
             if request.uuid != "":
                 # Filter by uuid
                 resource = await kb.get(request.uuid)
                 if resource:
-                    async for paragraph in resource.iterate_paragraphs(
-                        request.metadata
-                    ):
+                    async for paragraph in resource.iterate_paragraphs(request.metadata):
                         yield paragraph
             else:
                 async for resource in kb.iterate_resources():
-                    async for paragraph in resource.iterate_paragraphs(
-                        request.metadata
-                    ):
+                    async for paragraph in resource.iterate_paragraphs(request.metadata):
                         yield paragraph
 
     async def kb_fields(self, request: GetFieldsRequest) -> AsyncIterator[TrainField]:
@@ -129,15 +122,13 @@ class TrainShardManager(manager.KBShardManager):
                     async for field in resource.iterate_fields(request.metadata):
                         yield field
 
-    async def kb_resources(
-        self, request: GetResourcesRequest
-    ) -> AsyncIterator[TrainResource]:
+    async def kb_resources(self, request: GetResourcesRequest) -> AsyncIterator[TrainResource]:
         async with self.driver.transaction() as txn:
             kb = KnowledgeBox(txn, self.storage, request.kb.uuid)
             base = KB_RESOURCE_SLUG_BASE.format(kbid=request.kb.uuid)
-            async for key in txn.keys(match=base, count=-1):
+            async for key in txn.keys(match=base):
                 # Fetch and Add wanted item
-                rid = await txn.get(key)
+                rid = await txn.get(key, for_update=False)
                 if rid is not None:
                     resource = await kb.get(rid.decode())
                     if resource is not None:

nucliadb/train/servicer.py

@@ -18,10 +18,13 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import traceback
-from typing import Optional
 
 import aiohttp
-from nucliadb_protos.knowledgebox_pb2 import Labels
+
+from nucliadb.common import datamanagers
+from nucliadb.train.settings import settings
+from nucliadb.train.utils import get_shard_manager
+from nucliadb_protos import train_pb2_grpc
 from nucliadb_protos.train_pb2 import (
     GetFieldsRequest,
     GetInfoRequest,
@@ -38,10 +41,6 @@ from nucliadb_protos.writer_pb2 import (
     GetLabelsRequest,
     GetLabelsResponse,
 )
-
-from nucliadb.train.settings import settings
-from nucliadb.train.utils import get_shard_manager
-from nucliadb_protos import train_pb2_grpc
 from nucliadb_telemetry import errors
 
 
@@ -111,20 +110,15 @@ class TrainServicer(train_pb2_grpc.TrainServicer):
     async def GetOntology(  # type: ignore
         self, request: GetLabelsRequest, context=None
     ) -> GetLabelsResponse:
-        async with self.proc.driver.transaction() as txn:
-            kbobj = await self.proc.get_kb_obj(txn, request.kb.uuid)
-            labels: Optional[Labels] = None
-            if kbobj is not None:
-                labels = await kbobj.get_labels()
-
         response = GetLabelsResponse()
-        if kbobj is None:
-            response.status = GetLabelsResponse.Status.NOTFOUND
+        kbid = request.kb.uuid
+        labels = await datamanagers.atomic.labelset.get_all(kbid=kbid)
+        if labels is not None:
+            response.kb.uuid = kbid
+            response.status = GetLabelsResponse.Status.OK
+            response.labels.CopyFrom(labels)
         else:
-            response.kb.uuid = kbobj.kbid
-            if labels is not None:
-                response.labels.CopyFrom(labels)
-
+            response.status = GetLabelsResponse.Status.NOTFOUND
         return response
 
     async def GetOntologyCount(  # type: ignore
@@ -132,9 +126,7 @@ class TrainServicer(train_pb2_grpc.TrainServicer):
     ) -> LabelsetsCount:
         url = settings.internal_search_api.format(kbid=request.kb.uuid)
         facets = [f"faceted=/p/{labelset}" for labelset in request.paragraph_labelsets]
-        facets.extend(
-            [f"faceted=/l/{labelset}" for labelset in request.resource_labelsets]
-        )
+        facets.extend([f"faceted=/l/{labelset}" for labelset in request.resource_labelsets])
         query = "&".join(facets)
         headers = {"X-NUCLIADB-ROLES": "READER"}
         async with aiohttp.ClientSession() as sess:

nucliadb/train/settings.py

@@ -29,13 +29,9 @@ class Settings(DriverSettings):
     nuclia_learning_url: Optional[str] = "https://nuclia.cloud/api/v1/learning/"
     nuclia_learning_apikey: Optional[str] = None
 
-    internal_counter_api: str = (
-        "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/counters"
-    )
+    internal_counter_api: str = "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/counters"
 
-    internal_search_api: str = (
-        "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/search"
-    )
+    internal_search_api: str = "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/search"
 
 
 settings = Settings()

nucliadb/train/types.py

@@ -17,7 +17,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
-from typing import Union
+from typing import TypeVar, Union
 
 from nucliadb_protos import dataset_pb2 as dpb
 
@@ -29,14 +29,17 @@ TrainBatch = Union[
     dpb.QuestionAnswerStreamingBatch,
     dpb.SentenceClassificationBatch,
     dpb.TokenClassificationBatch,
+    dpb.FieldStreamingBatch,
 ]
 
-TrainBatchType = Union[
-    type[dpb.FieldClassificationBatch],
-    type[dpb.ImageClassificationBatch],
-    type[dpb.ParagraphClassificationBatch],
-    type[dpb.ParagraphStreamingBatch],
-    type[dpb.QuestionAnswerStreamingBatch],
-    type[dpb.SentenceClassificationBatch],
-    type[dpb.TokenClassificationBatch],
-]
+T = TypeVar(
+    "T",
+    dpb.FieldClassificationBatch,
+    dpb.ImageClassificationBatch,
+    dpb.ParagraphClassificationBatch,
+    dpb.ParagraphStreamingBatch,
+    dpb.QuestionAnswerStreamingBatch,
+    dpb.SentenceClassificationBatch,
+    dpb.TokenClassificationBatch,
+    dpb.FieldStreamingBatch,
+)

nucliadb/train/upload.py

@@ -19,11 +19,10 @@
 #
 import argparse
 import asyncio
+import importlib.metadata
 from asyncio import tasks
 from typing import Callable
 
-import pkg_resources
-
 from nucliadb.train.uploader import start_upload
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
@@ -33,9 +32,7 @@ from nucliadb_utils.settings import running_settings
 def arg_parse():
     parser = argparse.ArgumentParser(description="Upload data to Nuclia Learning API.")
 
-    parser.add_argument(
-        "-r", "--request", dest="request", help="Request UUID", required=True
-    )
+    parser.add_argument("-r", "--request", dest="request", help="Request UUID", required=True)
 
     parser.add_argument("-k", "--kb", dest="kb", help="Knowledge Box", required=True)
 
@@ -75,7 +72,7 @@ def _cancel_all_tasks(loop):
 def run() -> None:
     setup_logging()
 
-    errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
+    errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
     if asyncio._get_running_loop() is not None:
         raise RuntimeError("cannot be called from a running event loop")