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

nucliadb/tests/config.py

@@ -0,0 +1,53 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+
+# This is a test fixture which is useful outside nucliadb tests. In particular
+# it is used for the testbed. Keeping it under src so it can be imported outside
+def reset_config():
+    from nucliadb.common.cluster import settings as cluster_settings
+    from nucliadb.ingest import settings as ingest_settings
+    from nucliadb.train import settings as train_settings
+    from nucliadb.writer import settings as writer_settings
+    from nucliadb_utils import settings as utils_settings
+    from nucliadb_utils.cache import settings as cache_settings
+
+    all_settings = [
+        cluster_settings.settings,
+        ingest_settings.settings,
+        train_settings.settings,
+        writer_settings.settings,
+        cache_settings.settings,
+        utils_settings.audit_settings,
+        utils_settings.http_settings,
+        utils_settings.indexing_settings,
+        utils_settings.nuclia_settings,
+        utils_settings.nucliadb_settings,
+        utils_settings.storage_settings,
+        utils_settings.transaction_settings,
+    ]
+    for settings in all_settings:
+        defaults = type(settings)()
+        for attr, _value in settings:
+            default_value = getattr(defaults, attr)
+            setattr(settings, attr, default_value)
+
+    from nucliadb.common.cluster import manager
+
+    manager.INDEX_NODES.clear()

nucliadb/train/__init__.py

@@ -28,9 +28,7 @@ logger = logging.getLogger(SERVICE_NAME)
 class EndpointFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
         return (
-            record.args is not None
-            and len(record.args) >= 3
-            and record.args[2] not in ("/", "/metrics")  # type: ignore
+            record.args is not None and len(record.args) >= 3 and record.args[2] not in ("/", "/metrics")  # type: ignore
         )
 
 

nucliadb/train/api/utils.py

@@ -21,9 +21,8 @@
 
 from typing import Optional
 
-from nucliadb_protos.dataset_pb2 import TrainSet
-
 from nucliadb.train.utils import get_shard_manager
+from nucliadb_protos.dataset_pb2 import TrainSet
 
 
 async def get_kb_partitions(kbid: str, prefix: Optional[str] = None):

nucliadb/train/api/v1/shards.py

@@ -21,7 +21,7 @@
 
 from fastapi import HTTPException, Request
 from fastapi.responses import StreamingResponse
-from fastapi_versioning import version  # type: ignore
+from fastapi_versioning import version
 
 from nucliadb.train.api.utils import get_kb_partitions, get_train
 from nucliadb.train.api.v1.router import KB_PREFIX, api
@@ -34,7 +34,7 @@ from nucliadb_utils.authentication import requires_one
     f"/{KB_PREFIX}/{{kbid}}/trainset/{{shard}}",
     tags=["Object Response"],
     status_code=200,
-    name="Return Train Stream",
+    summary="Return Train Stream",
 )
 @requires_one([NucliaDBRoles.READER])
 @version(1)

nucliadb/train/api/v1/trainset.py

@@ -21,7 +21,7 @@
 from typing import Optional
 
 from fastapi import Request
-from fastapi_versioning import version  # type: ignore
+from fastapi_versioning import version
 
 from nucliadb.train.api.utils import get_kb_partitions
 from nucliadb.train.api.v1.router import KB_PREFIX, api
@@ -34,7 +34,7 @@ from nucliadb_utils.authentication import requires_one
     f"/{KB_PREFIX}/{{kbid}}/trainset",
     tags=["Train"],
     status_code=200,
-    name="Return Train call",
+    summary="Return Train call",
     response_model=TrainSetPartitions,
 )
 @requires_one([NucliaDBRoles.READER])
@@ -47,14 +47,12 @@ async def get_partitions_all(request: Request, kbid: str) -> TrainSetPartitions:
     f"/{KB_PREFIX}/{{kbid}}/trainset/{{prefix}}",
     tags=["Train"],
     status_code=200,
-    name="Return Train call",
+    summary="Return Train call",
     response_model=TrainSetPartitions,
 )
 @requires_one([NucliaDBRoles.READER])
 @version(1)
-async def get_partitions_prefix(
-    request: Request, kbid: str, prefix: str
-) -> TrainSetPartitions:
+async def get_partitions_prefix(request: Request, kbid: str, prefix: str) -> TrainSetPartitions:
     return await get_partitions(kbid, prefix=prefix)
 
 

nucliadb/train/app.py

@@ -17,77 +17,44 @@
 # 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/>.
 #
-import pkg_resources
+import importlib.metadata
+
 from fastapi import FastAPI
-from fastapi.responses import JSONResponse
 from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
-from starlette.middleware.cors import CORSMiddleware
 from starlette.requests import ClientDisconnect, Request
 from starlette.responses import HTMLResponse
 
-from nucliadb.middleware.transaction import ReadOnlyTransactionMiddleware
 from nucliadb.train import API_PREFIX
 from nucliadb.train.api.v1.router import api
-from nucliadb.train.lifecycle import finalize, initialize
+from nucliadb.train.lifecycle import lifespan
 from nucliadb_telemetry import errors
-from nucliadb_utils import const
+from nucliadb_telemetry.fastapi.utils import (
+    client_disconnect_handler,
+    global_exception_handler,
+)
+from nucliadb_utils.audit.stream import AuditMiddleware
 from nucliadb_utils.authentication import NucliaCloudAuthenticationBackend
 from nucliadb_utils.fastapi.openapi import extend_openapi
 from nucliadb_utils.fastapi.versioning import VersionedFastAPI
-from nucliadb_utils.settings import http_settings, running_settings
-from nucliadb_utils.utilities import has_feature
+from nucliadb_utils.settings import running_settings
+from nucliadb_utils.utilities import get_audit
 
 middleware = []
-
-if has_feature(const.Features.CORS_MIDDLEWARE, default=False):
-    middleware.append(
-        Middleware(
-            CORSMiddleware,
-            allow_origins=http_settings.cors_origins,
-            allow_methods=["*"],
-            # Authorization will be exluded from * in the future, (CORS non-wildcard request-header).
-            # Browsers already showing deprecation notices, so it needs to be specified explicitly
-            allow_headers=["*", "Authorization"],
-        )
-    )
-
 middleware.extend(
     [
-        Middleware(
-            AuthenticationMiddleware, backend=NucliaCloudAuthenticationBackend()
-        ),
-        Middleware(ReadOnlyTransactionMiddleware),
+        Middleware(AuthenticationMiddleware, backend=NucliaCloudAuthenticationBackend()),
+        Middleware(AuditMiddleware, audit_utility_getter=get_audit),
     ]
 )
 
-errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
-
-
-on_startup = [initialize]
-on_shutdown = [finalize]
-
-
-async def global_exception_handler(request: Request, exc: Exception):
-    errors.capture_exception(exc)
-    return JSONResponse(
-        status_code=500,
-        content={"detail": "Something went wrong, please contact your administrator"},
-    )
-
-
-async def client_disconnect_handler(request: Request, exc: ClientDisconnect):
-    return JSONResponse(
-        status_code=200,
-        content={"detail": "Client disconnected while an operation was in course"},
-    )
+errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
 
 fastapi_settings = dict(
     debug=running_settings.debug,
     middleware=middleware,
-    on_startup=on_startup,
-    on_shutdown=on_shutdown,
+    lifespan=lifespan,
     exception_handlers={
         Exception: global_exception_handler,
         ClientDisconnect: client_disconnect_handler,

nucliadb/train/generator.py

@@ -21,11 +21,11 @@
 from typing import AsyncIterator, Optional
 
 from fastapi import HTTPException
-from nucliadb_protos.dataset_pb2 import TaskType, TrainSet
 
 from nucliadb.train.generators.field_classifier import (
     field_classification_batch_generator,
 )
+from nucliadb.train.generators.field_streaming import field_streaming_batch_generator
 from nucliadb.train.generators.image_classifier import (
     image_classification_batch_generator,
 )
@@ -46,6 +46,7 @@ from nucliadb.train.generators.token_classifier import (
 )
 from nucliadb.train.types import TrainBatch
 from nucliadb.train.utils import get_shard_manager
+from nucliadb_protos.dataset_pb2 import TaskType, TrainSet
 
 
 async def generate_train_data(kbid: str, shard: str, trainset: TrainSet):
@@ -59,34 +60,24 @@ async def generate_train_data(kbid: str, shard: str, trainset: TrainSet):
     batch_generator: Optional[AsyncIterator[TrainBatch]] = None
 
     if trainset.type == TaskType.FIELD_CLASSIFICATION:
-        batch_generator = field_classification_batch_generator(
-            kbid, trainset, node, shard_replica_id
-        )
+        batch_generator = field_classification_batch_generator(kbid, trainset, node, shard_replica_id)
     elif trainset.type == TaskType.IMAGE_CLASSIFICATION:
-        batch_generator = image_classification_batch_generator(
-            kbid, trainset, node, shard_replica_id
-        )
+        batch_generator = image_classification_batch_generator(kbid, trainset, node, shard_replica_id)
     elif trainset.type == TaskType.PARAGRAPH_CLASSIFICATION:
         batch_generator = paragraph_classification_batch_generator(
             kbid, trainset, node, shard_replica_id
         )
     elif trainset.type == TaskType.TOKEN_CLASSIFICATION:
-        batch_generator = token_classification_batch_generator(
-            kbid, trainset, node, shard_replica_id
-        )
+        batch_generator = token_classification_batch_generator(kbid, trainset, node, shard_replica_id)
     elif trainset.type == TaskType.SENTENCE_CLASSIFICATION:
-        batch_generator = sentence_classification_batch_generator(
-            kbid, trainset, node, shard_replica_id
-        )
+        batch_generator = sentence_classification_batch_generator(kbid, trainset, node, shard_replica_id)
     elif trainset.type == TaskType.PARAGRAPH_STREAMING:
-        batch_generator = paragraph_streaming_batch_generator(
-            kbid, trainset, node, shard_replica_id
-        )
+        batch_generator = paragraph_streaming_batch_generator(kbid, trainset, node, shard_replica_id)
 
     elif trainset.type == TaskType.QUESTION_ANSWER_STREAMING:
-        batch_generator = question_answer_batch_generator(
-            kbid, trainset, node, shard_replica_id
-        )
+        batch_generator = question_answer_batch_generator(kbid, trainset, node, shard_replica_id)
+    elif trainset.type == TaskType.FIELD_STREAMING:
+        batch_generator = field_streaming_batch_generator(kbid, trainset, node, shard_replica_id)
 
     if batch_generator is None:
         raise HTTPException(

nucliadb/train/generators/field_classifier.py

@@ -20,7 +20,10 @@
 
 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 (
     FieldClassificationBatch,
     Label,
@@ -29,11 +32,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 field_classification_batch_generator(
     kbid: str,
@@ -41,15 +39,7 @@ def field_classification_batch_generator(
     node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[FieldClassificationBatch, None]:
-    if len(trainset.filter.labels) != 1:
-        raise HTTPException(
-            status_code=422,
-            detail="Paragraph Classification should be of 1 labelset",
-        )
-
-    generator = generate_field_classification_payloads(
-        kbid, trainset, node, shard_replica_id
-    )
+    generator = generate_field_classification_payloads(kbid, trainset, node, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, FieldClassificationBatch)
     return batch_generator
 
@@ -95,13 +85,11 @@ async def get_field_text(kbid: str, rid: str, field: str, field_type: str) -> st
         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 ""
 
     text = ""

nucliadb/train/generators/field_streaming.py

@@ -0,0 +1,156 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 AsyncGenerator, Optional
+
+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 (
+    FieldSplitData,
+    FieldStreamingBatch,
+    TrainSet,
+)
+from nucliadb_protos.nodereader_pb2 import StreamRequest
+from nucliadb_protos.resources_pb2 import Basic, FieldComputedMetadata
+from nucliadb_protos.utils_pb2 import ExtractedText
+
+
+def field_streaming_batch_generator(
+    kbid: str,
+    trainset: TrainSet,
+    node: AbstractIndexNode,
+    shard_replica_id: str,
+) -> AsyncGenerator[FieldStreamingBatch, None]:
+    generator = generate_field_streaming_payloads(kbid, trainset, node, shard_replica_id)
+    batch_generator = batchify(generator, trainset.batch_size, FieldStreamingBatch)
+    return batch_generator
+
+
+async def generate_field_streaming_payloads(
+    kbid: str,
+    trainset: TrainSet,
+    node: AbstractIndexNode,
+    shard_replica_id: str,
+) -> AsyncGenerator[FieldSplitData, None]:
+    # Query how many resources has each label
+    request = StreamRequest()
+    request.shard_id.id = shard_replica_id
+
+    for label in trainset.filter.labels:
+        request.filter.labels.append(f"/l/{label}")
+
+    for path in trainset.filter.paths:
+        request.filter.labels.append(f"/p/{path}")
+
+    for metadata in trainset.filter.metadata:
+        request.filter.labels.append(f"/m/{metadata}")
+
+    for entity in trainset.filter.entities:
+        request.filter.labels.append(f"/e/{entity}")
+
+    for field in trainset.filter.fields:
+        request.filter.labels.append(f"/f/{field}")
+
+    for status in trainset.filter.status:
+        request.filter.labels.append(f"/n/s/{status}")
+    total = 0
+
+    async for document_item in node.stream_get_fields(request):
+        text_labels = []
+        for label in document_item.labels:
+            text_labels.append(label)
+
+        field_id = f"{document_item.uuid}{document_item.field}"
+        total += 1
+
+        field_parts = document_item.field.split("/")
+        if len(field_parts) == 3:
+            _, field_type, field = field_parts
+            split = "0"
+        elif len(field_parts) == 4:
+            _, field_type, field, split = field_parts
+        else:
+            raise Exception(f"Invalid field definition {document_item.field}")
+
+        tl = FieldSplitData()
+        rid, field_type, field = field_id.split("/")
+        tl.rid = document_item.uuid
+        tl.field = field
+        tl.field_type = field_type
+        tl.split = split
+        extracted = await get_field_text(kbid, rid, field, field_type)
+        if extracted is not None:
+            tl.text.CopyFrom(extracted)
+
+        metadata_obj = await get_field_metadata(kbid, rid, field, field_type)
+        if metadata_obj is not None:
+            tl.metadata.CopyFrom(metadata_obj)
+
+        basic = await get_field_basic(kbid, rid, field, field_type)
+        if basic is not None:
+            tl.basic.CopyFrom(basic)
+
+        tl.labels.extend(text_labels)
+
+        yield tl
+
+
+async def get_field_text(kbid: str, rid: str, field: str, field_type: str) -> Optional[ExtractedText]:
+    orm_resource = await get_resource_from_cache_or_db(kbid, rid)
+
+    if orm_resource is None:
+        logger.error(f"{rid} does not exist on DB")
+        return None
+
+    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()
+
+    return extracted_text
+
+
+async def get_field_metadata(
+    kbid: str, rid: str, field: str, field_type: str
+) -> Optional[FieldComputedMetadata]:
+    orm_resource = await get_resource_from_cache_or_db(kbid, rid)
+
+    if orm_resource is None:
+        logger.error(f"{rid} does not exist on DB")
+        return None
+
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
+    field_obj = await orm_resource.get_field(field, field_type_int, load=False)
+    field_metadata = await field_obj.get_field_metadata()
+
+    return field_metadata
+
+
+async def get_field_basic(kbid: str, rid: str, field: str, field_type: str) -> Optional[Basic]:
+    orm_resource = await get_resource_from_cache_or_db(kbid, rid)
+
+    if orm_resource is None:
+        logger.error(f"{rid} does not exist on DB")
+        return None
+
+    basic = await orm_resource.get_basic()
+
+    return basic

nucliadb/train/generators/image_classifier.py

@@ -21,6 +21,12 @@
 import json
 from typing import Any, AsyncGenerator
 
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.ingest.fields.base import Field
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.train import logger
+from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
     ImageClassification,
     ImageClassificationBatch,
@@ -29,12 +35,6 @@ from nucliadb_protos.dataset_pb2 import (
 from nucliadb_protos.nodereader_pb2 import StreamRequest
 from nucliadb_protos.resources_pb2 import FieldType, PageStructure, VisualSelection
 
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.ingest.fields.base import Field
-from nucliadb.ingest.orm.resource import KB_REVERSE, Resource
-from nucliadb.train import logger
-from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
-
 VISUALLY_ANNOTABLE_FIELDS = {FieldType.FILE, FieldType.LINK}
 
 # PAWLS JSON format
@@ -47,9 +47,7 @@ def image_classification_batch_generator(
     node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[ImageClassificationBatch, None]:
-    generator = generate_image_classification_payloads(
-        kbid, trainset, node, shard_replica_id
-    )
+    generator = generate_image_classification_payloads(kbid, trainset, node, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, ImageClassificationBatch)
     return batch_generator
 
@@ -71,7 +69,7 @@ async def generate_image_classification_payloads(
             return
 
         _, field_type_key, field_key = item.field.split("/")
-        field_type = KB_REVERSE[field_type_key]
+        field_type = FIELD_TYPE_STR_TO_PB[field_type_key]
 
         if field_type not in VISUALLY_ANNOTABLE_FIELDS:
             continue
@@ -131,9 +129,7 @@ async def generate_image_classification_payloads(
             yield ic
 
 
-async def get_page_selections(
-    resource: Resource, field: Field
-) -> dict[int, list[VisualSelection]]:
+async def get_page_selections(resource: Resource, field: Field) -> dict[int, list[VisualSelection]]:
     page_selections: dict[int, list[VisualSelection]] = {}
     basic = await resource.get_basic()
     if basic is None or basic.fieldmetadata is None:
@@ -144,7 +140,7 @@ async def get_page_selections(
     for fieldmetadata in basic.fieldmetadata:
         if (
             fieldmetadata.field.field == field.id
-            and fieldmetadata.field.field_type == KB_REVERSE[field.type]
+            and fieldmetadata.field.field_type == FIELD_TYPE_STR_TO_PB[field.type]
         ):
             for selection in fieldmetadata.page_selections:
                 page_selections[selection.page] = selection.visual  # type: ignore
@@ -155,7 +151,7 @@ async def get_page_selections(
 
 async def get_page_structure(field: Field) -> list[tuple[str, PageStructure]]:
     page_structures: list[tuple[str, PageStructure]] = []
-    field_type = KB_REVERSE[field.type]
+    field_type = FIELD_TYPE_STR_TO_PB[field.type]
     if field_type == FieldType.FILE:
         fed = await field.get_file_extracted_data()  # type: ignore
         if fed is None:
@@ -163,9 +159,7 @@ async def get_page_structure(field: Field) -> list[tuple[str, PageStructure]]:
 
         fp = fed.file_pages_previews
         if len(fp.pages) != len(fp.structures):
-            field_path = (
-                f"/kb/{field.kbid}/resource/{field.resource.uuid}/file/{field.id}"
-            )
+            field_path = f"/kb/{field.kbid}/resource/{field.resource.uuid}/file/{field.id}"
             logger.warning(
                 f"File extracted data has a different number of pages and structures! ({field_path})"
             )

nucliadb/train/generators/paragraph_classifier.py

@@ -21,6 +21,9 @@
 from typing import AsyncGenerator
 
 from fastapi import HTTPException
+
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.train.generators.utils import batchify, get_paragraph
 from nucliadb_protos.dataset_pb2 import (
     Label,
     ParagraphClassificationBatch,
@@ -29,9 +32,6 @@ from nucliadb_protos.dataset_pb2 import (
 )
 from nucliadb_protos.nodereader_pb2 import StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.train.generators.utils import batchify, get_paragraph
-
 
 def paragraph_classification_batch_generator(
     kbid: str,
@@ -45,12 +45,8 @@ def paragraph_classification_batch_generator(
             detail="Paragraph Classification should be of 1 labelset",
         )
 
-    generator = generate_paragraph_classification_payloads(
-        kbid, trainset, node, shard_replica_id
-    )
-    batch_generator = batchify(
-        generator, trainset.batch_size, ParagraphClassificationBatch
-    )
+    generator = generate_paragraph_classification_payloads(kbid, trainset, node, shard_replica_id)
+    batch_generator = batchify(generator, trainset.batch_size, ParagraphClassificationBatch)
     return batch_generator
 
 

nucliadb/train/generators/paragraph_streaming.py

@@ -20,6 +20,10 @@
 
 from typing import AsyncGenerator
 
+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 (
     ParagraphStreamingBatch,
     ParagraphStreamItem,
@@ -27,11 +31,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 paragraph_streaming_batch_generator(
     kbid: str,
@@ -39,9 +38,7 @@ def paragraph_streaming_batch_generator(
     node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[ParagraphStreamingBatch, None]:
-    generator = generate_paragraph_streaming_payloads(
-        kbid, trainset, node, shard_replica_id
-    )
+    generator = generate_paragraph_streaming_payloads(kbid, trainset, node, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, ParagraphStreamingBatch)
     return batch_generator
 
@@ -68,7 +65,7 @@ async def generate_paragraph_streaming_payloads(
             logger.error(f"{rid} does not exist on DB")
             continue
 
-        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()