doccano/backend/data_export/pipeline/repositories.py

import abc
import itertools
from collections import defaultdict
from typing import Any, Dict, Iterator, List, Tuple

from .data import Record
from examples.models import Example
from projects.models import Project

SpanType = Tuple[int, int, str]


class BaseRepository:
    def __init__(self, project: Project):
        self.project = project

    def list(self, export_approved=False) -> Iterator[Record]:
        raise NotImplementedError()

    def create_unlabeled_record(self, example: Example) -> Record:
        raise NotImplementedError()


class FileRepository(BaseRepository):
    def list(self, export_approved=False) -> Iterator[Record]:
        examples = self.project.examples.all()
        if export_approved:
            examples = examples.exclude(states=None)

        for example in examples:
            label_per_user = self.label_per_user(example)
            if self.project.collaborative_annotation:
                label_per_user = self.reduce_user(label_per_user)
            for user, label in label_per_user.items():
                yield Record(
                    data_id=example.id,
                    data=example.upload_name,
                    label=label,
                    user=user,
                    metadata=example.meta,
                )
            # todo:
            # If there is no label, export the doc with `unknown` user.
            # This is a quick solution.
            # In the future, the doc without label will be exported
            # with the user who approved the doc.
            # This means I will allow each user to be able to approve the doc.
            if len(label_per_user) == 0:
                yield self.create_unlabeled_record(example)

    def create_unlabeled_record(self, example: Example) -> Record:
        return Record(data_id=example.id, data=example.upload_name, label=[], user="unknown", metadata=example.meta)

    def label_per_user(self, example) -> Dict:
        label_per_user = defaultdict(list)
        for a in example.categories.all():
            label_per_user[a.user.username].append(a.label.text)
        return label_per_user

    def reduce_user(self, label_per_user: Dict[str, Any]):
        value = list(itertools.chain(*label_per_user.values()))
        return {"all": value}


class Speech2TextRepository(FileRepository):
    def label_per_user(self, example) -> Dict:
        label_per_user = defaultdict(list)
        for a in example.texts.all():
            label_per_user[a.user.username].append(a.text)
        return label_per_user


class TextRepository(BaseRepository):
    @property
    def docs(self):
        return Example.objects.filter(project=self.project)

    def list(self, export_approved=False):
        docs = self.docs
        if export_approved:
            docs = docs.exclude(states=None)

        for doc in docs:
            label_per_user = self.label_per_user(doc)
            if self.project.collaborative_annotation:
                label_per_user = self.reduce_user(label_per_user)
            for user, label in label_per_user.items():
                yield Record(data_id=doc.id, data=doc.text, label=label, user=user, metadata=doc.meta)
            # todo:
            # If there is no label, export the doc with `unknown` user.
            # This is a quick solution.
            # In the future, the doc without label will be exported
            # with the user who approved the doc.
            # This means I will allow each user to be able to approve the doc.
            if len(label_per_user) == 0:
                yield self.create_unlabeled_record(doc)

    def create_unlabeled_record(self, example: Example) -> Record:
        return Record(data_id=example.id, data=example.text, label=[], user="unknown", metadata=example.meta)

    @abc.abstractmethod
    def label_per_user(self, doc) -> Dict:
        raise NotImplementedError()

    def reduce_user(self, label_per_user: Dict[str, Any]):
        value = list(itertools.chain(*label_per_user.values()))
        return {"all": value}


class TextClassificationRepository(TextRepository):
    @property
    def docs(self):
        return Example.objects.filter(project=self.project).prefetch_related("categories__user", "categories__label")

    def label_per_user(self, doc) -> Dict:
        label_per_user = defaultdict(list)
        for a in doc.categories.all():
            label_per_user[a.user.username].append(a.label.text)
        return label_per_user


class SequenceLabelingRepository(TextRepository):
    @property
    def docs(self):
        return Example.objects.filter(project=self.project).prefetch_related("spans__user", "spans__label")

    def label_per_user(self, doc) -> Dict:
        label_per_user = defaultdict(list)
        for a in doc.spans.all():
            label = (a.start_offset, a.end_offset, a.label.text)
            label_per_user[a.user.username].append(label)
        return label_per_user


class RelationExtractionRepository(TextRepository):
    @property
    def docs(self):
        return Example.objects.filter(project=self.project).prefetch_related(
            "spans__user", "spans__label", "relations__user", "relations__type"
        )

    def create_unlabeled_record(self, example: Example) -> Record:
        return Record(
            data_id=example.id,
            data=example.text,
            label={"entities": [], "relations": []},
            user="unknown",
            metadata=example.meta,
        )

    def label_per_user(self, doc) -> Dict:
        relation_per_user: Dict = defaultdict(list)
        span_per_user: Dict = defaultdict(list)
        label_per_user: Dict = defaultdict(dict)
        for relation in doc.relations.all():
            relation_per_user[relation.user.username].append(
                {
                    "id": relation.id,
                    "from_id": relation.from_id.id,
                    "to_id": relation.to_id.id,
                    "type": relation.type.text,
                }
            )
        for span in doc.spans.all():
            span_per_user[span.user.username].append(
                {
                    "id": span.id,
                    "start_offset": span.start_offset,
                    "end_offset": span.end_offset,
                    "label": span.label.text,
                }
            )
        for user, relations in relation_per_user.items():
            label_per_user[user]["relations"] = relations
        for user, span in span_per_user.items():
            label_per_user[user]["entities"] = span
        return label_per_user

    def reduce_user(self, label_per_user: Dict[str, Any]):
        entities = []
        relations = []
        for user, label in label_per_user.items():
            entities.extend(label.get("entities", []))
            relations.extend(label.get("relations", []))
        return {"all": {"entities": entities, "relations": relations}}


class Seq2seqRepository(TextRepository):
    @property
    def docs(self):
        return Example.objects.filter(project=self.project).prefetch_related("texts__user")

    def label_per_user(self, doc) -> Dict:
        label_per_user = defaultdict(list)
        for a in doc.texts.all():
            label_per_user[a.user.username].append(a.text)
        return label_per_user


class IntentDetectionSlotFillingRepository(TextRepository):
    @property
    def docs(self):
        return Example.objects.filter(project=self.project).prefetch_related(
            "categories__user", "categories__label", "spans__user", "spans__label"
        )

    def create_unlabeled_record(self, example: Example) -> Record:
        return Record(
            data_id=example.id,
            data=example.text,
            label={"entities": [], "cats": []},
            user="unknown",
            metadata=example.meta,
        )

    def label_per_user(self, doc) -> Dict:
        category_per_user: Dict[str, List[str]] = defaultdict(list)
        span_per_user: Dict[str, List[SpanType]] = defaultdict(list)
        label_per_user: Dict[str, Dict[str, List]] = defaultdict(dict)
        for a in doc.categories.all():
            category_per_user[a.user.username].append(a.label.text)
        for a in doc.spans.all():
            span_per_user[a.user.username].append((a.start_offset, a.end_offset, a.label.text))
        for user, cats in category_per_user.items():
            label_per_user[user]["cats"] = cats
        for user, span in span_per_user.items():
            label_per_user[user]["entities"] = span
        for label in label_per_user.values():
            label.setdefault("cats", [])
            label.setdefault("entities", [])
        return label_per_user

    def reduce_user(self, label_per_user: Dict[str, Any]):
        cats = []
        entities = []
        for user, label in label_per_user.items():
            cats.extend(label.get("cats", []))
            entities.extend(label.get("entities", []))
        return {"all": {"entities": entities, "cats": cats}}