nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.whl

nucliadb/train/tests/test_token_classification.py

@@ -1,136 +0,0 @@
-# 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/>.
-
-import asyncio
-
-import aiohttp
-import pytest
-from nucliadb_protos.dataset_pb2 import TaskType, TokenClassificationBatch, TrainSet
-from nucliadb_protos.resources_pb2 import Position
-from nucliadb_protos.writer_pb2 import BrokerMessage
-from nucliadb_protos.writer_pb2_grpc import WriterStub
-
-from nucliadb.tests.utils import inject_message
-from nucliadb.tests.utils.broker_messages import BrokerMessageBuilder, FieldBuilder
-from nucliadb.train import API_PREFIX
-from nucliadb.train.api.v1.router import KB_PREFIX
-from nucliadb.train.tests.utils import get_batches_from_train_response_stream
-from nucliadb_protos import resources_pb2 as rpb
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("knowledgebox", ["STABLE", "EXPERIMENTAL"], indirect=True)
-async def test_generator_token_classification(
-    train_rest_api: aiohttp.ClientSession,
-    knowledgebox_with_entities: str,
-    nucliadb_grpc: WriterStub,
-):
-    kbid = knowledgebox_with_entities
-
-    await inject_resource_with_token_classification(kbid, nucliadb_grpc)
-
-    async with train_rest_api.get(
-        f"/{API_PREFIX}/v1/{KB_PREFIX}/{kbid}/trainset"
-    ) as partitions:
-        assert partitions.status == 200
-        data = await partitions.json()
-        assert len(data["partitions"]) == 1
-        partition_id = data["partitions"][0]
-
-    trainset = TrainSet()
-    trainset.type = TaskType.TOKEN_CLASSIFICATION
-    trainset.batch_size = 2
-    trainset.filter.labels.append("PERSON")
-    trainset.filter.labels.append("ORG")
-    async with train_rest_api.post(
-        f"/{API_PREFIX}/v1/{KB_PREFIX}/{kbid}/trainset/{partition_id}",
-        data=trainset.SerializeToString(),
-    ) as response:
-        assert response.status == 200
-        batches: list[TokenClassificationBatch] = []
-        async for batch in get_batches_from_train_response_stream(
-            response, TokenClassificationBatch
-        ):
-            batches.append(batch)
-
-    for batch in batches:
-        if batch.data[0].token == "Eudald":
-            assert batch.data[0].label == "B-PERSON"
-            assert batch.data[1].label == "I-PERSON"
-            assert batch.data[2].label == "O"
-        if batch.data[0].token == "This":
-            assert batch.data[4].label == "B-PERSON"
-            assert batch.data[5].label == "I-PERSON"
-        if batch.data[0].token == "Where":
-            assert batch.data[3].label == "B-ORG"
-            assert batch.data[4].label == "I-ORG"
-            assert batch.data[5].label == "I-ORG"
-        if batch.data[0].token == "Summary":
-            assert batch.data[2].label == "B-ORG"
-            assert batch.data[4].label == "B-ORG"
-        if batch.data[0].token == "My":
-            assert batch.data[3].label == "B-PERSON"
-            assert batch.data[12].label == "B-ORG"
-
-
-async def inject_resource_with_token_classification(knowledgebox, writer):
-    bm = broker_resource(knowledgebox)
-    await inject_message(writer, bm)
-    await asyncio.sleep(0.1)
-    return bm.uuid
-
-
-def broker_resource(knowledgebox: str) -> BrokerMessage:
-    bmb = BrokerMessageBuilder(kbid=knowledgebox)
-
-    bmb.with_title("This is a bird, its a plane, no, its el Super Fran")
-    title_field = bmb.field_builder("title", rpb.FieldType.GENERIC)
-    title_field.with_extracted_entity(
-        "PERSON", "el Super Fran", positions=[Position(start=37, end=50)]
-    )
-
-    bmb.with_summary("Summary of Nuclia using Debian")
-    summary_field = bmb.field_builder("summary", rpb.FieldType.GENERIC)
-    summary_field.with_extracted_entity(
-        "ORG", "Nuclia", positions=[Position(start=11, end=17)]
-    )
-    summary_field.with_extracted_entity(
-        "ORG", "Debian", positions=[Position(start=24, end=30)]
-    )
-
-    file_field = FieldBuilder("file", rpb.FieldType.FILE)
-    file_field.with_extracted_text(
-        "My own text Ramon. This is great to be at Nuclia. \n Where is the Generalitat de Catalunya? Eudald Camprubi, do you want to go shooping? This is a test Carmen Iniesta!"  # noqa
-    )
-    file_field.with_extracted_paragraph_metadata(rpb.Paragraph(start=0, end=49))
-    file_field.with_extracted_paragraph_metadata(rpb.Paragraph(start=50, end=90))
-    file_field.with_extracted_paragraph_metadata(rpb.Paragraph(start=91, end=135))
-    file_field.with_extracted_paragraph_metadata(rpb.Paragraph(start=136, end=166))
-
-    file_field.with_user_entity("PERSON", "Ramon", start=12, end=17)
-    file_field.with_user_entity("ORG", "Nuclia", start=42, end=48)
-    file_field.with_user_entity("ORG", "Generalitat de Catalunya", start=65, end=89)
-    file_field.with_user_entity("PERSON", "Eudald", start=91, end=106)
-    file_field.with_user_entity("PERSON", "Carmen Iniesta", start=151, end=165)
-
-    bmb.add_field_builder(file_field)
-
-    bm = bmb.build()
-
-    return bm

nucliadb/train/tests/utils.py

@@ -1,101 +0,0 @@
-# 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, overload
-
-import aiohttp
-from nucliadb_protos.dataset_pb2 import (
-    FieldClassificationBatch,
-    ImageClassificationBatch,
-    ParagraphClassificationBatch,
-    ParagraphStreamingBatch,
-    QuestionAnswerStreamingBatch,
-    SentenceClassificationBatch,
-    TokenClassificationBatch,
-)
-
-from nucliadb.train.types import TrainBatch, TrainBatchType
-
-# NOTE: we use def instead of async def to make mypy happy. Otherwise, it
-# considers the overloaded functions as corountines returning async iterators
-# instead of async iterators themselves and complains about it
-
-
-@overload
-def get_batches_from_train_response_stream(
-    response: aiohttp.ClientResponse,
-    pb_klass: type[FieldClassificationBatch],
-) -> AsyncGenerator[FieldClassificationBatch, None]: ...
-
-
-@overload
-def get_batches_from_train_response_stream(
-    response: aiohttp.ClientResponse,
-    pb_klass: type[ImageClassificationBatch],
-) -> AsyncGenerator[ImageClassificationBatch, None]: ...
-
-
-@overload
-def get_batches_from_train_response_stream(
-    response: aiohttp.ClientResponse,
-    pb_klass: type[ParagraphClassificationBatch],
-) -> AsyncGenerator[ParagraphClassificationBatch, None]: ...
-
-
-@overload
-def get_batches_from_train_response_stream(
-    response: aiohttp.ClientResponse,
-    pb_klass: type[ParagraphStreamingBatch],
-) -> AsyncGenerator[ParagraphStreamingBatch, None]: ...
-
-
-@overload
-def get_batches_from_train_response_stream(
-    response: aiohttp.ClientResponse,
-    pb_klass: type[QuestionAnswerStreamingBatch],
-) -> AsyncGenerator[QuestionAnswerStreamingBatch, None]: ...
-
-
-@overload
-def get_batches_from_train_response_stream(
-    response: aiohttp.ClientResponse,
-    pb_klass: type[SentenceClassificationBatch],
-) -> AsyncGenerator[SentenceClassificationBatch, None]: ...
-
-
-@overload
-def get_batches_from_train_response_stream(
-    response: aiohttp.ClientResponse,
-    pb_klass: type[TokenClassificationBatch],
-) -> AsyncGenerator[TokenClassificationBatch, None]: ...
-
-
-async def get_batches_from_train_response_stream(
-    response: aiohttp.ClientResponse,
-    pb_klass: TrainBatchType,
-) -> AsyncGenerator[TrainBatch, None]:
-    while True:
-        header = await response.content.read(4)
-        if header == b"":
-            break
-        payload_size = int.from_bytes(header, byteorder="big", signed=False)
-        payload = await response.content.read(payload_size)
-        batch = pb_klass()
-        batch.ParseFromString(payload)
-        yield batch

nucliadb/writer/layouts/__init__.py

@@ -1,51 +0,0 @@
-# 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 Any, Callable, Coroutine
-
-from nucliadb_protos.resources_pb2 import FieldLayout
-
-import nucliadb_models as models
-from nucliadb_utils.storages.storage import Storage
-
-VERSION: dict[
-    int,
-    Callable[
-        [models.InputLayoutField, str, str, str, Storage],
-        Coroutine[Any, Any, FieldLayout],
-    ],
-] = {}
-
-import nucliadb.writer.layouts.v1  # noqa isort:skip
-
-
-async def serialize_blocks(
-    layout_field: models.InputLayoutField,
-    kbid: str,
-    uuid: str,
-    field: str,
-    storage: Storage,
-) -> FieldLayout:
-    if layout_field.format in VERSION:
-        layout = await VERSION[layout_field.format](
-            layout_field, kbid, uuid, field, storage
-        )
-    else:
-        raise KeyError("Invalid version")
-    return layout

nucliadb/writer/layouts/v1.py

@@ -1,59 +0,0 @@
-# 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 nucliadb_protos.resources_pb2 import Block as PBBlock
-from nucliadb_protos.resources_pb2 import FieldLayout
-
-import nucliadb_models as models
-from nucliadb.writer.layouts import VERSION
-from nucliadb_utils.storages.storage import Storage
-
-
-async def serialize_block(
-    layout_field: models.InputLayoutField,
-    kbid: str,
-    uuid: str,
-    field: str,
-    storage: Storage,
-) -> FieldLayout:
-    pblayout = FieldLayout()
-    for key, block in layout_field.body.blocks.items():
-        pbblock = PBBlock()
-        pbblock.x = block.x
-        pbblock.y = block.y
-        pbblock.cols = block.cols
-        pbblock.rows = block.rows
-        pbblock.type = PBBlock.TypeBlock.Value(block.type)
-        pbblock.ident = block.ident if block.ident else key
-        pbblock.payload = block.payload
-
-        sf = storage.layout_field(kbid, uuid, field, key)
-        await storage.upload_b64file_to_cloudfile(
-            sf,
-            block.file.payload.encode(),
-            block.file.filename,
-            block.file.content_type,
-            block.file.md5,
-        )
-        pblayout.body.blocks[key].CopyFrom(pbblock)
-        pblayout.format = FieldLayout.Format.Value(layout_field.format.value)
-    return pblayout
-
-
-VERSION[models.LayoutFormat.NUCLIAv1] = serialize_block

nucliadb/writer/tests/__init__.py

@@ -1,19 +0,0 @@
-# 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/>.
-#

nucliadb/writer/tests/conftest.py

@@ -1,31 +0,0 @@
-# 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/>.
-#
-pytest_plugins = [
-    "pytest_mock",
-    "pytest_docker_fixtures",
-    "nucliadb_utils.tests.nats",
-    "nucliadb.tests.fixtures",
-    "nucliadb.tests.tikv",
-    "nucliadb.ingest.tests.fixtures",  # should be refactored out
-    "nucliadb.writer.tests.fixtures",
-    "nucliadb_utils.tests.conftest",
-    "nucliadb_utils.tests.gcs",
-    "nucliadb_utils.tests.s3",
-]

nucliadb/writer/tests/fixtures.py

@@ -1,191 +0,0 @@
-# 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 enum import Enum
-from typing import AsyncIterator, Callable, Optional
-from unittest import mock
-
-import pytest
-from httpx import AsyncClient
-from pytest_lazy_fixtures import lazy_fixture
-from redis import asyncio as aioredis
-
-from nucliadb.ingest.tests.fixtures import IngestFixture
-from nucliadb.writer import API_PREFIX
-from nucliadb.writer.api.v1.router import KB_PREFIX, KBS_PREFIX
-from nucliadb.writer.settings import settings
-from nucliadb.writer.tus import clear_storage
-from nucliadb_models.resource import NucliaDBRoles
-from nucliadb_utils.settings import (
-    FileBackendConfig,
-    nuclia_settings,
-    nucliadb_settings,
-    storage_settings,
-)
-from nucliadb_utils.tests.conftest import get_testing_storage_backend
-from nucliadb_utils.utilities import Utility, clean_utility, set_utility
-
-
-@pytest.fixture(scope="function")
-def disabled_back_pressure():
-    with mock.patch(
-        "nucliadb.writer.back_pressure.is_back_pressure_enabled", return_value=False
-    ) as mocked:
-        yield mocked
-
-
-@pytest.fixture(scope="function")
-async def writer_api(
-    disabled_back_pressure,
-    redis,
-    storage_writer,
-    grpc_servicer: IngestFixture,
-    transaction_utility,
-    processing_utility,
-    tus_manager,
-) -> AsyncIterator[Callable[[list[Enum], str, str], AsyncClient]]:
-    nucliadb_settings.nucliadb_ingest = grpc_servicer.host
-    from nucliadb.writer.app import create_application
-
-    application = create_application()
-
-    def make_client_fixture(
-        roles: Optional[list[Enum]] = None,
-        user: str = "",
-        version: str = "1",
-    ) -> AsyncClient:
-        roles = roles or []
-        client_base_url = "http://test"
-        client_base_url = f"{client_base_url}/{API_PREFIX}/v{version}"
-
-        client = AsyncClient(app=application, base_url=client_base_url)  # type: ignore
-        client.headers["X-NUCLIADB-ROLES"] = ";".join(
-            map(lambda role: role.value, roles)
-        )
-        client.headers["X-NUCLIADB-USER"] = user
-
-        return client
-
-    driver = aioredis.from_url(f"redis://{redis[0]}:{redis[1]}")
-    await driver.flushall()
-
-    await application.router.startup()
-
-    yield make_client_fixture
-
-    await application.router.shutdown()
-    clear_storage()
-
-    await driver.flushall()
-    await driver.close(close_connection_pool=True)
-
-
-@pytest.fixture(scope="function")
-def gcs_storage_writer(gcs):
-    storage_settings.file_backend = FileBackendConfig.GCS
-    storage_settings.gcs_endpoint_url = gcs
-    storage_settings.gcs_bucket = "test_{kbid}"
-
-
-@pytest.fixture(scope="function")
-def s3_storage_writer(s3):
-    storage_settings.file_backend = FileBackendConfig.S3
-    storage_settings.s3_endpoint = s3
-    storage_settings.s3_client_id = ""
-    storage_settings.s3_client_secret = ""
-    storage_settings.s3_bucket = "test-{kbid}"
-
-
-@pytest.fixture(scope="function")
-def pg_storage_writer(pg):
-    storage_settings.file_backend = FileBackendConfig.PG
-    url = f"postgresql://postgres:postgres@{pg[0]}:{pg[1]}/postgres"
-    storage_settings.driver_pg_url = url
-
-
-def lazy_storage_writer_fixture():
-    backend = get_testing_storage_backend()
-    if backend == "gcs":
-        return [lazy_fixture.lf("gcs_storage_writer")]
-    elif backend == "s3":
-        return [lazy_fixture.lf("s3_storage_writer")]
-    elif backend == "pg":
-        return [lazy_fixture.lf("pg_storage_writer")]
-    else:
-        print(f"Unknown storage backend {backend}, using gcs")
-        return [lazy_fixture.lf("gcs_storage_writer")]
-
-
-@pytest.fixture(scope="function", params=lazy_storage_writer_fixture())
-async def storage_writer(request):
-    """
-    Generic storage fixture that allows us to run the same tests for different storage backends.
-    """
-    storage_driver = request.param
-    set_utility(Utility.STORAGE, storage_driver)
-
-    yield storage_driver
-
-    clean_utility(Utility.STORAGE)
-
-
-@pytest.fixture(scope="function")
-async def knowledgebox_writer(writer_api):
-    async with writer_api(roles=[NucliaDBRoles.MANAGER]) as client:
-        resp = await client.post(
-            f"/{KBS_PREFIX}",
-            json={
-                "slug": "kbid1",
-                "title": "My Knowledge Box",
-            },
-        )
-        assert resp.status_code == 201
-    kbid = resp.json().get("uuid")
-    assert kbid is not None
-    yield kbid
-
-
-@pytest.fixture(scope="function")
-async def resource(redis, writer_api, knowledgebox_writer):
-    async with writer_api(roles=[NucliaDBRoles.WRITER]) as client:
-        resp = await client.post(
-            f"/{KB_PREFIX}/{knowledgebox_writer}/resources",
-            json={
-                "slug": "resource1",
-                "title": "Resource 1",
-            },
-        )
-        assert resp.status_code == 201
-        uuid = resp.json()["uuid"]
-
-    return uuid
-
-
-@pytest.fixture(scope="function")
-async def processing_utility():
-    nuclia_settings.dummy_processing = True
-    nuclia_settings.onprem = True
-    nuclia_settings.nuclia_jwt_key = "foobarkey"
-
-
-@pytest.fixture(scope="function")
-async def tus_manager(redis):
-    settings.dm_redis_host = redis[0]
-    settings.dm_redis_port = redis[1]
-    yield