nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__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,108 +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/resource/vectors.py

@@ -1,120 +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 Optional
-
-from fastapi import HTTPException
-from nucliadb_protos.resources_pb2 import UserVectorsWrapper
-from nucliadb_protos.writer_pb2 import (
-    BrokerMessage,
-    GetVectorSetsRequest,
-    GetVectorSetsResponse,
-    OpStatusWriter,
-    SetVectorSetRequest,
-    VectorSets,
-)
-
-from nucliadb_models.common import FIELD_TYPES_MAP_REVERSE
-from nucliadb_models.vectors import UserVectorsWrapper as UserVectorsWrapperPy
-from nucliadb_models.vectors import VectorSimilarity
-from nucliadb_utils.utilities import get_ingest
-
-
-async def create_vectorset(
-    kbid: str,
-    vectorset: str,
-    dimension: Optional[int] = None,
-    similarity: Optional[VectorSimilarity] = None,
-):
-    ingest = get_ingest()
-    pbrequest: SetVectorSetRequest = SetVectorSetRequest(id=vectorset)
-    pbrequest.kb.uuid = kbid
-
-    if dimension is not None:
-        pbrequest.vectorset.dimension = dimension
-    if similarity:
-        pbrequest.vectorset.similarity = similarity.to_pb()
-
-    status: OpStatusWriter = await ingest.SetVectorSet(pbrequest)  # type: ignore
-    if status.status == OpStatusWriter.Status.OK:
-        return None
-    elif status.status == OpStatusWriter.Status.NOTFOUND:
-        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    elif status.status == OpStatusWriter.Status.ERROR:
-        raise HTTPException(
-            status_code=500, detail="Error on settings labels on a Knowledge box"
-        )
-
-
-async def get_vectorsets(kbid: str) -> Optional[VectorSets]:
-    ingest = get_ingest()
-    pbrequest: GetVectorSetsRequest = GetVectorSetsRequest()
-    pbrequest.kb.uuid = kbid
-
-    vectorsets: GetVectorSetsResponse = await ingest.GetVectorSets(pbrequest)  # type: ignore
-    if vectorsets.status == GetVectorSetsResponse.Status.OK:
-        return vectorsets.vectorsets
-    else:
-        return None
-
-
-def parse_vectors(
-    writer: BrokerMessage, vectors: UserVectorsWrapperPy, vectorsets: VectorSets
-):
-    for vector in vectors:
-        evw = UserVectorsWrapper()
-        evw.field.field_type = FIELD_TYPES_MAP_REVERSE[
-            vector.field.field_type.value
-        ]  # type: ignore
-        evw.field.field = vector.field.field
-        if vector.vectors is not None:
-            for vectorset, user_vectors in vector.vectors.items():
-                if vectorset not in vectorsets.vectorsets:
-                    raise HTTPException(
-                        status_code=412,
-                        detail=str(f"Invalid vectorset"),
-                    )
-                else:
-                    dimension = vectorsets.vectorsets[vectorset].dimension
-                    for key, user_vector in user_vectors.items():
-                        vo = evw.vectors.vectors[vectorset].vectors[key]
-                        if len(user_vector.vector) == dimension:
-                            if user_vector.vector.count(0) == dimension:
-                                raise HTTPException(
-                                    status_code=412,
-                                    detail=str(f"Invalid vector should not be 0"),
-                                )
-                            vo.vector.extend(user_vector.vector)
-                        else:
-                            raise HTTPException(
-                                status_code=412,
-                                detail=str(
-                                    f"Invalid dimension should be {dimension} was {len(user_vector.vector)}"
-                                ),
-                            )
-                        if user_vector.positions is not None:
-                            vo.start = user_vector.positions[0]
-                            vo.end = user_vector.positions[1]
-
-        if vector.vectors_to_delete is not None:
-            for vectorset, user_vector_list in vector.vectors_to_delete.items():
-                evw.vectors_to_delete[vectorset].vectors.extend(
-                    user_vector_list.vectors
-                )
-
-        writer.user_vectors.append(evw)

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",
-]