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

nucliadb/ingest/orm/brain.py

@@ -22,17 +22,25 @@ from copy import deepcopy
 from dataclasses import dataclass
 from typing import Optional
 
+from nucliadb.common import ids
+from nucliadb.ingest import logger
+from nucliadb.ingest.orm.utils import compute_paragraph_key
+from nucliadb_models.labels import BASE_LABELS, LABEL_HIDDEN, flatten_resource_labels
+from nucliadb_models.metadata import ResourceProcessingStatus
+from nucliadb_protos import utils_pb2
 from nucliadb_protos.noderesources_pb2 import IndexParagraph as BrainParagraph
-from nucliadb_protos.noderesources_pb2 import ParagraphMetadata
+from nucliadb_protos.noderesources_pb2 import (
+    ParagraphMetadata,
+    Representation,
+    ResourceID,
+)
 from nucliadb_protos.noderesources_pb2 import Position as TextPosition
-from nucliadb_protos.noderesources_pb2 import Representation
 from nucliadb_protos.noderesources_pb2 import Resource as PBBrainResource
-from nucliadb_protos.noderesources_pb2 import ResourceID
 from nucliadb_protos.resources_pb2 import (
     Basic,
     ExtractedText,
+    FieldAuthor,
     FieldComputedMetadata,
-    FieldKeywordset,
     FieldMetadata,
     Metadata,
     Origin,
@@ -40,25 +48,10 @@ from nucliadb_protos.resources_pb2 import (
     UserFieldMetadata,
     UserMetadata,
 )
-from nucliadb_protos.utils_pb2 import Relation, RelationNode, VectorObject
-
-from nucliadb.ingest import logger
-from nucliadb.ingest.orm.utils import compute_paragraph_key
-from nucliadb_models.labels import BASE_LABELS, flatten_resource_labels
-from nucliadb_models.metadata import ResourceProcessingStatus
-from nucliadb_protos import utils_pb2
+from nucliadb_protos.utils_pb2 import Relation, RelationNode
 
 FilePagePositions = dict[int, tuple[int, int]]
 
-FIELD_PARAGRAPH_ID = "{rid}/{field_id}/{paragraph_start}-{paragraph_end}"
-SPLIT_FIELD_PARAGRAPH_ID = (
-    "{rid}/{field_id}/{subfield_id}/{paragraph_start}-{paragraph_end}"
-)
-FIELD_VECTOR_ID = "{rid}/{field_id}/{index}/{vector_start}-{vector_end}"
-SPLIT_FIELD_VECTOR_ID = (
-    "{rid}/{field_id}/{subfield_id}/{index}/{vector_start}-{vector_end}"
-)
-
 METADATA_STATUS_PB_TYPE_TO_NAME_MAP = {
     Metadata.Status.ERROR: ResourceProcessingStatus.ERROR.name,
     Metadata.Status.PROCESSED: ResourceProcessingStatus.PROCESSED.name,
@@ -79,7 +72,7 @@ class ResourceBrain:
         self.rid = rid
         ridobj = ResourceID(uuid=rid)
         self.brain: PBBrainResource = PBBrainResource(resource=ridobj)
-        self.labels: dict[str, list[str]] = deepcopy(BASE_LABELS)
+        self.labels: dict[str, set[str]] = deepcopy(BASE_LABELS)
 
     def apply_field_text(self, field_key: str, text: str):
         self.brain.texts[field_key].text = text
@@ -104,28 +97,30 @@ class ResourceBrain:
         self,
         field_key: str,
         metadata: FieldComputedMetadata,
-        replace_field: list[str],
-        replace_splits: dict[str, list[str]],
         page_positions: Optional[FilePagePositions],
         extracted_text: Optional[ExtractedText],
         basic_user_field_metadata: Optional[UserFieldMetadata] = None,
+        *,
+        replace_field: bool = False,
     ):
         # To check for duplicate paragraphs
         unique_paragraphs: set[str] = set()
 
         # Expose also user classifications
-        paragraph_classifications = self._get_paragraph_user_classifications(
+        user_paragraph_classifications = self._get_paragraph_user_classifications(
             basic_user_field_metadata
         )
 
         # We should set paragraphs and labels
         paragraph_pages = ParagraphPages(page_positions) if page_positions else None
         for subfield, metadata_split in metadata.split_metadata.items():
+            extracted_text_str = extracted_text.split_text[subfield] if extracted_text else None
+
             # For each split of this field
             for index, paragraph in enumerate(metadata_split.paragraphs):
                 key = f"{self.rid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
 
-                denied_classifications = paragraph_classifications.denied.get(key, [])
+                denied_classifications = set(user_paragraph_classifications.denied.get(key, []))
                 position = TextPosition(
                     index=index,
                     start=paragraph.start,
@@ -157,9 +152,8 @@ class ResourceBrain:
                     index=index,
                     repeated_in_field=is_paragraph_repeated_in_field(
                         paragraph,
-                        extracted_text,
+                        extracted_text_str,
                         unique_paragraphs,
-                        split=subfield,
                     ),
                     metadata=ParagraphMetadata(
                         position=position,
@@ -167,22 +161,22 @@ class ResourceBrain:
                         representation=representation,
                     ),
                 )
-                p.labels.append(
-                    f"/k/{Paragraph.TypeParagraph.Name(paragraph.kind).lower()}"
+                paragraph_kind_label = f"/k/{Paragraph.TypeParagraph.Name(paragraph.kind).lower()}"
+                paragraph_labels = {paragraph_kind_label}
+                paragraph_labels.update(
+                    f"/l/{classification.labelset}/{classification.label}"
+                    for classification in paragraph.classifications
                 )
-                for classification in paragraph.classifications:
-                    label = f"/l/{classification.labelset}/{classification.label}"
-                    if label not in denied_classifications:
-                        p.labels.append(label)
-
-                # Add user annotated labels to paragraphs
-                extend_unique(p.labels, paragraph_classifications.valid.get(key, []))  # type: ignore
+                paragraph_labels.update(set(user_paragraph_classifications.valid.get(key, [])))
+                paragraph_labels.difference_update(denied_classifications)
+                p.labels.extend(list(paragraph_labels))
 
                 self.brain.paragraphs[field_key].paragraphs[key].CopyFrom(p)
 
+        extracted_text_str = extracted_text.text if extracted_text else None
         for index, paragraph in enumerate(metadata.metadata.paragraphs):
             key = f"{self.rid}/{field_key}/{paragraph.start}-{paragraph.end}"
-            denied_classifications = paragraph_classifications.denied.get(key, [])
+            denied_classifications = set(user_paragraph_classifications.denied.get(key, []))
             position = TextPosition(
                 index=index,
                 start=paragraph.start,
@@ -212,7 +206,7 @@ class ResourceBrain:
                 field=field_key,
                 index=index,
                 repeated_in_field=is_paragraph_repeated_in_field(
-                    paragraph, extracted_text, unique_paragraphs
+                    paragraph, extracted_text_str, unique_paragraphs
                 ),
                 metadata=ParagraphMetadata(
                     position=position,
@@ -220,72 +214,59 @@ class ResourceBrain:
                     representation=representation,
                 ),
             )
-            p.labels.append(
-                f"/k/{Paragraph.TypeParagraph.Name(paragraph.kind).lower()}"
+            paragraph_kind_label = f"/k/{Paragraph.TypeParagraph.Name(paragraph.kind).lower()}"
+            paragraph_labels = {paragraph_kind_label}
+            paragraph_labels.update(
+                f"/l/{classification.labelset}/{classification.label}"
+                for classification in paragraph.classifications
             )
-
-            for classification in paragraph.classifications:
-                label = f"/l/{classification.labelset}/{classification.label}"
-                if label not in denied_classifications:
-                    p.labels.append(label)
-
-            # Add user annotated labels to paragraphs
-            extend_unique(p.labels, paragraph_classifications.valid.get(key, []))  # type: ignore
+            paragraph_labels.update(set(user_paragraph_classifications.valid.get(key, [])))
+            paragraph_labels.difference_update(denied_classifications)
+            p.labels.extend(list(paragraph_labels))
 
             self.brain.paragraphs[field_key].paragraphs[key].CopyFrom(p)
 
+        if replace_field:
+            field_type, field_name = field_key.split("/")
+            full_field_id = ids.FieldId(rid=self.rid, type=field_type, key=field_name).full()
+            self.brain.paragraphs_to_delete.append(full_field_id)
+
         for relations in metadata.metadata.relations:
             for relation in relations.relations:
                 self.brain.relations.append(relation)
 
-        for split, sentences in replace_splits.items():
-            for sentence in sentences:
-                self.brain.paragraphs_to_delete.append(
-                    f"{self.rid}/{field_key}/{split}/{sentence}"
-                )
-
-        for sentence_to_delete in replace_field:
-            self.brain.paragraphs_to_delete.append(
-                f"{self.rid}/{field_key}/{sentence_to_delete}"
-            )
-
-    def delete_metadata(self, field_key: str, metadata: FieldComputedMetadata):
-        for subfield, metadata_split in metadata.split_metadata.items():
-            for paragraph in metadata_split.paragraphs:
-                self.brain.paragraphs_to_delete.append(
-                    f"{self.rid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
-                )
-
-        for paragraph in metadata.metadata.paragraphs:
-            self.brain.sentences_to_delete.append(
-                f"{self.rid}/{field_key}/{paragraph.start}-{paragraph.end}"
-            )
+    def delete_field(self, field_key: str):
+        ftype, fkey = field_key.split("/")
+        full_field_id = ids.FieldId(rid=self.rid, type=ftype, key=fkey).full()
+        self.brain.paragraphs_to_delete.append(full_field_id)
+        self.brain.sentences_to_delete.append(full_field_id)
 
     def apply_field_vectors(
         self,
         field_id: str,
         vo: utils_pb2.VectorObject,
         *,
+        vectorset: Optional[str] = None,
         replace_field: bool = False,
-        replace_splits: Optional[list[str]] = None,
         matryoshka_vector_dimension: Optional[int] = None,
     ):
-        replace_splits = replace_splits or []
-
+        fid = ids.FieldId.from_string(f"{self.rid}/{field_id}")
         for subfield, vectors in vo.split_vectors.items():
+            _field_id = ids.FieldId(
+                rid=fid.rid,
+                type=fid.type,
+                key=fid.key,
+                subfield_id=subfield,
+            )
             # For each split of this field
             for index, vector in enumerate(vectors.vectors):
-                paragraph_key = SPLIT_FIELD_PARAGRAPH_ID.format(
-                    rid=self.rid,
-                    field_id=field_id,
-                    subfield_id=subfield,
+                paragraph_key = ids.ParagraphId(
+                    field_id=_field_id,
                     paragraph_start=vector.start_paragraph,
                     paragraph_end=vector.end_paragraph,
                 )
-                sentence_key = SPLIT_FIELD_VECTOR_ID.format(
-                    rid=self.rid,
-                    field_id=field_id,
-                    subfield_id=subfield,
+                sentence_key = ids.VectorId(
+                    field_id=_field_id,
                     index=index,
                     vector_start=vector.start,
                     vector_end=vector.end,
@@ -295,19 +276,23 @@ class ResourceBrain:
                     paragraph_key,
                     sentence_key,
                     vector,
+                    vectorset=vectorset,
                     matryoshka_vector_dimension=matryoshka_vector_dimension,
                 )
 
+        _field_id = ids.FieldId(
+            rid=fid.rid,
+            type=fid.type,
+            key=fid.key,
+        )
         for index, vector in enumerate(vo.vectors.vectors):
-            paragraph_key = FIELD_PARAGRAPH_ID.format(
-                rid=self.rid,
-                field_id=field_id,
+            paragraph_key = ids.ParagraphId(
+                field_id=_field_id,
                 paragraph_start=vector.start_paragraph,
                 paragraph_end=vector.end_paragraph,
             )
-            sentence_key = FIELD_VECTOR_ID.format(
-                rid=self.rid,
-                field_id=field_id,
+            sentence_key = ids.VectorId(
+                field_id=_field_id,
                 index=index,
                 vector_start=vector.start,
                 vector_end=vector.end,
@@ -317,26 +302,33 @@ class ResourceBrain:
                 paragraph_key,
                 sentence_key,
                 vector,
+                vectorset=vectorset,
                 matryoshka_vector_dimension=matryoshka_vector_dimension,
             )
 
-        for split in replace_splits:
-            self.brain.sentences_to_delete.append(f"{self.rid}/{field_id}/{split}")
-
         if replace_field:
-            self.brain.sentences_to_delete.append(f"{self.rid}/{field_id}")
+            full_field_id = ids.FieldId(rid=self.rid, type=fid.type, key=fid.key).full()
+            if vectorset is None:
+                # DEPRECATED
+                self.brain.sentences_to_delete.append(full_field_id)
+            else:
+                self.brain.vector_prefixes_to_delete[vectorset].items.append(full_field_id)
 
     def _apply_field_vector(
         self,
         field_id: str,
-        paragraph_key: str,
-        sentence_key: str,
+        paragraph_key: ids.ParagraphId,
+        sentence_key: ids.VectorId,
         vector: utils_pb2.Vector,
         *,
+        vectorset: Optional[str],
         matryoshka_vector_dimension: Optional[int] = None,
     ):
-        paragraph_pb = self.brain.paragraphs[field_id].paragraphs[paragraph_key]
-        sentence_pb = paragraph_pb.sentences[sentence_key]
+        paragraph_pb = self.brain.paragraphs[field_id].paragraphs[paragraph_key.full()]
+        if vectorset:
+            sentence_pb = paragraph_pb.vectorsets_sentences[vectorset].sentences[sentence_key.full()]
+        else:
+            sentence_pb = paragraph_pb.sentences[sentence_key.full()]
 
         sentence_pb.ClearField("vector")  # clear first to prevent duplicates
 
@@ -352,39 +344,18 @@ class ResourceBrain:
         sentence_pb.metadata.position.end = vector.end_paragraph
 
         # does it make sense to copy forward paragraph values here?
-        sentence_pb.metadata.position.page_number = (
-            paragraph_pb.metadata.position.page_number
-        )
+        sentence_pb.metadata.position.page_number = paragraph_pb.metadata.position.page_number
         sentence_pb.metadata.position.in_page = paragraph_pb.metadata.position.in_page
 
         sentence_pb.metadata.page_with_visual = paragraph_pb.metadata.page_with_visual
 
-        sentence_pb.metadata.representation.file = (
-            paragraph_pb.metadata.representation.file
-        )
+        sentence_pb.metadata.representation.file = paragraph_pb.metadata.representation.file
 
-        sentence_pb.metadata.representation.is_a_table = (
-            paragraph_pb.metadata.representation.is_a_table
-        )
+        sentence_pb.metadata.representation.is_a_table = paragraph_pb.metadata.representation.is_a_table
 
         sentence_pb.metadata.position.index = paragraph_pb.metadata.position.index
 
-    def delete_vectors(self, field_key: str, vo: VectorObject):
-        # TODO: no need to iterate over all vectors, just delete the whole field
-        for subfield, vectors in vo.split_vectors.items():
-            for vector in vectors.vectors:
-                self.brain.sentences_to_delete.append(
-                    f"{self.rid}/{field_key}/{subfield}/{vector.start}-{vector.end}"
-                )
-
-        for vector in vo.vectors.vectors:
-            self.brain.sentences_to_delete.append(
-                f"{self.rid}/{field_key}/{vector.start}-{vector.end}"
-            )
-
-    def set_processing_status(
-        self, basic: Basic, previous_status: Optional[Metadata.Status.ValueType]
-    ):
+    def set_processing_status(self, basic: Basic, previous_status: Optional[Metadata.Status.ValueType]):
         """
         We purposefully overwrite what we index as a status and DO NOT reflect
         actual status with what we index.
@@ -441,15 +412,11 @@ class ResourceBrain:
                 self.brain.metadata.modified.CopyFrom(origin.modified)
 
     def _set_resource_relations(self, basic: Basic, origin: Optional[Origin]):
-        relationnodedocument = RelationNode(
-            value=self.rid, ntype=RelationNode.NodeType.RESOURCE
-        )
+        relationnodedocument = RelationNode(value=self.rid, ntype=RelationNode.NodeType.RESOURCE)
         if origin is not None:
             # origin contributors
             for contrib in origin.colaborators:
-                relationnodeuser = RelationNode(
-                    value=contrib, ntype=RelationNode.NodeType.USER
-                )
+                relationnodeuser = RelationNode(value=contrib, ntype=RelationNode.NodeType.USER)
                 self.brain.relations.append(
                     Relation(
                         relation=Relation.COLAB,
@@ -478,115 +445,147 @@ class ResourceBrain:
     def _set_resource_labels(self, basic: Basic, origin: Optional[Origin]):
         if origin is not None:
             if origin.source_id:
-                self.labels["o"] = [origin.source_id]
+                self.labels["o"] = {origin.source_id}
             # origin tags
             for tag in origin.tags:
-                self.labels["t"].append(tag)
+                self.labels["t"].add(tag)
             # origin source
             if origin.source_id != "":
-                self.labels["u"].append(f"s/{origin.source_id}")
+                self.labels["u"].add(f"s/{origin.source_id}")
 
             if origin.path:
-                self.labels["p"].append(origin.path.lstrip("/"))
+                self.labels["p"].add(origin.path.lstrip("/"))
 
             # origin contributors
             for contrib in origin.colaborators:
-                self.labels["u"].append(f"o/{contrib}")
+                self.labels["u"].add(f"o/{contrib}")
 
             for key, value in origin.metadata.items():
-                self.labels["m"].append(f"{key[:255]}/{value[:255]}")
+                self.labels["m"].add(f"{key[:255]}/{value[:255]}")
 
         # icon
-        self.labels["n"].append(f"i/{basic.icon}")
+        self.labels["n"].add(f"i/{basic.icon}")
 
         # processing status
         status_tag = self.get_processing_status_tag(basic.metadata)
-        self.labels["n"].append(f"s/{status_tag}")
+        self.labels["n"].add(f"s/{status_tag}")
 
         # main language
         if basic.metadata.language:
-            self.labels["s"].append(f"p/{basic.metadata.language}")
+            self.labels["s"].add(f"p/{basic.metadata.language}")
 
         # all language
         for lang in basic.metadata.languages:
-            self.labels["s"].append(f"s/{lang}")
+            self.labels["s"].add(f"s/{lang}")
 
         # labels
         for classification in basic.usermetadata.classifications:
-            self.labels["l"].append(f"{classification.labelset}/{classification.label}")
+            self.labels["l"].add(f"{classification.labelset}/{classification.label}")
 
-        self.compute_labels()
+        # hidden
+        if basic.hidden:
+            _, p1, p2 = LABEL_HIDDEN.split("/")
+            self.labels[p1].add(p2)
+
+        self.brain.ClearField("labels")
+        self.brain.labels.extend(flatten_resource_labels(self.labels))
 
     def process_field_metadata(
         self,
         field_key: str,
         metadata: FieldMetadata,
-        labels: dict[str, list[str]],
+        labels: dict[str, set[str]],
         relation_node_document: RelationNode,
-        user_canceled_labels: list[str],
+        user_canceled_labels: set[str],
     ):
+        if metadata.mime_type != "":
+            labels["mt"].add(metadata.mime_type)
+
+        base_classification_relation = Relation(
+            relation=Relation.ABOUT,
+            source=relation_node_document,
+            to=RelationNode(
+                ntype=RelationNode.NodeType.LABEL,
+            ),
+        )
         for classification in metadata.classifications:
             label = f"{classification.labelset}/{classification.label}"
             if label not in user_canceled_labels:
-                labels["l"].append(label)
-                relation_node_label = RelationNode(
-                    value=label,
-                    ntype=RelationNode.NodeType.LABEL,
-                )
-                self.brain.relations.append(
-                    Relation(
-                        relation=Relation.ABOUT,
-                        source=relation_node_document,
-                        to=relation_node_label,
-                    )
-                )
+                labels["l"].add(label)
+                relation = Relation()
+                relation.CopyFrom(base_classification_relation)
+                relation.to.value = label
+                self.brain.relations.append(relation)
+
+        # Data Augmentation + Processor entities
+        base_entity_relation = Relation(
+            relation=Relation.ENTITY,
+            source=relation_node_document,
+            to=RelationNode(ntype=RelationNode.NodeType.ENTITY),
+        )
+        use_legacy_entities = True
+        for data_augmentation_task_id, entities in metadata.entities.items():
+            # If we recieved the entities from the processor here, we don't want to use the legacy entities
+            # TODO: Remove this when processor doesn't use this anymore
+            if data_augmentation_task_id == "processor":
+                use_legacy_entities = False
+
+            for ent in entities.entities:
+                entity_text = ent.text
+                entity_label = ent.label
+                # Seems like we don't care about where the entity is in the text
+                # entity_positions = entity.positions
+                labels["e"].add(
+                    f"{entity_label}/{entity_text}"
+                )  # Add data_augmentation_task_id as a prefix?
+                relation = Relation()
+                relation.CopyFrom(base_entity_relation)
+                relation.to.value = entity_text
+                relation.to.subtype = entity_label
+                self.brain.relations.append(relation)
 
-        for klass_entity, _ in metadata.positions.items():
-            labels["e"].append(klass_entity)
-            entity_array = klass_entity.split("/")
-            if len(entity_array) == 1:
+        # Legacy processor entities
+        # TODO: Remove once processor doesn't use this anymore and remove the positions and ner fields from the message
+        def _parse_entity(klass_entity: str) -> tuple[str, str]:
+            try:
+                klass, entity = klass_entity.split("/", 1)
+                return klass, entity
+            except ValueError:
                 raise AttributeError(f"Entity should be with type {klass_entity}")
-            elif len(entity_array) > 1:
-                klass = entity_array[0]
-                entity = "/".join(entity_array[1:])
-            relation_node_entity = RelationNode(
-                value=entity, ntype=RelationNode.NodeType.ENTITY, subtype=klass
-            )
-            rel = Relation(
-                relation=Relation.ENTITY,
-                source=relation_node_document,
-                to=relation_node_entity,
-            )
-            self.brain.relations.append(rel)
 
-    def process_keywordset_fields(self, field_key: str, field: FieldKeywordset):
-        # all field keywords
-        if field:
-            for keyword in field.keywords:
-                self.labels["f"].append(f"{field_key}/{keyword.value}")
-                self.labels["fg"].append(keyword.value)
+        if use_legacy_entities:
+            for klass_entity in metadata.positions.keys():
+                labels["e"].add(klass_entity)
+                klass, entity = _parse_entity(klass_entity)
+                relation = Relation()
+                relation.CopyFrom(base_entity_relation)
+                relation.to.value = entity
+                relation.to.subtype = klass
+                self.brain.relations.append(relation)
 
     def apply_field_labels(
         self,
         field_key: str,
         metadata: Optional[FieldComputedMetadata],
         uuid: str,
+        generated_by: FieldAuthor,
         basic_user_metadata: Optional[UserMetadata] = None,
         basic_user_fieldmetadata: Optional[UserFieldMetadata] = None,
     ):
+        user_canceled_labels: set[str] = set()
         if basic_user_metadata is not None:
-            user_canceled_labels = [
-                f"/l/{classification.labelset}/{classification.label}"
+            user_canceled_labels.update(
+                f"{classification.labelset}/{classification.label}"
                 for classification in basic_user_metadata.classifications
                 if classification.cancelled_by_user
-            ]
-        else:
-            user_canceled_labels = []
-
-        relation_node_resource = RelationNode(
-            value=uuid, ntype=RelationNode.NodeType.RESOURCE
-        )
-        labels: dict[str, list[str]] = {"l": [], "e": []}
+            )
+        relation_node_resource = RelationNode(value=uuid, ntype=RelationNode.NodeType.RESOURCE)
+        labels: dict[str, set[str]] = {
+            "l": set(),  # classification labels
+            "e": set(),  # entities
+            "mt": set(),  # mime type
+            "g/da": set(),  # generated by
+        }
         if metadata is not None:
             for meta in metadata.split_metadata.values():
                 self.process_field_metadata(
@@ -607,7 +606,7 @@ class ResourceBrain:
         if basic_user_fieldmetadata is not None:
             for token in basic_user_fieldmetadata.token:
                 if token.cancelled_by_user is False:
-                    labels["e"].append(f"{token.klass}/{token.token}")
+                    labels["e"].add(f"{token.klass}/{token.token}")
                     relation_node_entity = RelationNode(
                         value=token.token,
                         ntype=RelationNode.NodeType.ENTITY,
@@ -635,36 +634,33 @@ class ResourceBrain:
                             self.brain.paragraphs[field_key].paragraphs[
                                 paragraph_annotation.key
                             ].labels.append(label)
-        extend_unique(
-            self.brain.texts[field_key].labels, flatten_resource_labels(labels)  # type: ignore
-        )
-
-    def compute_labels(self):
-        extend_unique(self.brain.labels, flatten_resource_labels(self.labels))
 
+        if generated_by.WhichOneof("author") == "data_augmentation":
+            field_type, field_id = field_key.split("/")
+            da_task_id = ids.extract_data_augmentation_id(field_id)
+            if da_task_id is None:  # pragma: nocover
+                logger.warning(
+                    "Data augmentation field id has an unexpected format! Skipping label",
+                    extra={
+                        "rid": uuid,
+                        "field_id": field_id,
+                    },
+                )
+            else:
+                labels["g/da"].add(da_task_id)
 
-def get_paragraph_text(
-    extracted_text: ExtractedText, start: int, end: int, split: Optional[str] = None
-) -> str:
-    if split is not None:
-        text = extracted_text.split_text[split]
-    else:
-        text = extracted_text.text
-    return text[start:end]
+        self.brain.texts[field_key].labels.extend(flatten_resource_labels(labels))
 
 
 def is_paragraph_repeated_in_field(
     paragraph: Paragraph,
-    extracted_text: Optional[ExtractedText],
+    extracted_text: Optional[str],
     unique_paragraphs: set[str],
-    split: Optional[str] = None,
 ) -> bool:
     if extracted_text is None:
         return False
 
-    paragraph_text = get_paragraph_text(
-        extracted_text, start=paragraph.start, end=paragraph.end, split=split
-    )
+    paragraph_text = extracted_text[paragraph.start : paragraph.end]
     if len(paragraph_text) == 0:
         return False
 
@@ -701,12 +697,3 @@ class ParagraphPages:
             if len(self._materialized) > 0:
                 return self._materialized[-1]
             return 0
-
-
-def extend_unique(a: list, b: list):
-    """
-    Prevents extending with duplicate elements
-    """
-    for item in b:
-        if item not in a:
-            a.append(item)