arkindex-base-worker 0.4.0rc1__tar.gz → 0.4.0rc2__tar.gz

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.4.0rc1
+Version: 0.4.0rc2
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>
@@ -40,7 +40,7 @@ Classifier: Programming Language :: Python :: 3.11
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: humanize==4.9.0
+Requires-Dist: humanize==4.10.0
 Requires-Dist: peewee~=3.17
 Requires-Dist: Pillow==10.4.0
 Requires-Dist: python-gnupg==0.5.2
@@ -49,7 +49,7 @@ Requires-Dist: teklia-toolbox==0.1.5
 Requires-Dist: zstandard==0.22.0
 Provides-Extra: docs
 Requires-Dist: black==24.4.2; extra == "docs"
-Requires-Dist: mkdocs-material==9.5.31; extra == "docs"
+Requires-Dist: mkdocs-material==9.5.33; extra == "docs"
 Requires-Dist: mkdocstrings-python==1.10.8; extra == "docs"
 Provides-Extra: tests
 Requires-Dist: pytest==8.3.2; extra == "tests"

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/arkindex_base_worker.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.4.0rc1
+Version: 0.4.0rc2
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>
@@ -40,7 +40,7 @@ Classifier: Programming Language :: Python :: 3.11
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: humanize==4.9.0
+Requires-Dist: humanize==4.10.0
 Requires-Dist: peewee~=3.17
 Requires-Dist: Pillow==10.4.0
 Requires-Dist: python-gnupg==0.5.2
@@ -49,7 +49,7 @@ Requires-Dist: teklia-toolbox==0.1.5
 Requires-Dist: zstandard==0.22.0
 Provides-Extra: docs
 Requires-Dist: black==24.4.2; extra == "docs"
-Requires-Dist: mkdocs-material==9.5.31; extra == "docs"
+Requires-Dist: mkdocs-material==9.5.33; extra == "docs"
 Requires-Dist: mkdocstrings-python==1.10.8; extra == "docs"
 Provides-Extra: tests
 Requires-Dist: pytest==8.3.2; extra == "tests"

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/arkindex_base_worker.egg-info/SOURCES.txt

@@ -20,6 +20,7 @@ arkindex_worker/worker/element.py
 arkindex_worker/worker/entity.py
 arkindex_worker/worker/image.py
 arkindex_worker/worker/metadata.py
+arkindex_worker/worker/process.py
 arkindex_worker/worker/task.py
 arkindex_worker/worker/training.py
 arkindex_worker/worker/transcription.py

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/arkindex_base_worker.egg-info/requires.txt

@@ -1,4 +1,4 @@
-humanize==4.9.0
+humanize==4.10.0
 peewee~=3.17
 Pillow==10.4.0
 python-gnupg==0.5.2
@@ -8,7 +8,7 @@ zstandard==0.22.0
 
 [docs]
 black==24.4.2
-mkdocs-material==9.5.31
+mkdocs-material==9.5.33
 mkdocstrings-python==1.10.8
 
 [tests]

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/arkindex_worker/worker/__init__.py

@@ -4,12 +4,10 @@ Base classes to implement Arkindex workers.
 
 import contextlib
 import json
-import os
 import sys
 import uuid
-from argparse import ArgumentTypeError
-from collections.abc import Iterable, Iterator
-from enum import Enum
+from collections.abc import Iterable
+from itertools import chain
 from pathlib import Path
 
 from apistar.exceptions import ErrorResponse
@@ -21,47 +19,27 @@ from arkindex_worker.utils import pluralize
 from arkindex_worker.worker.base import BaseWorker
 from arkindex_worker.worker.classification import ClassificationMixin
 from arkindex_worker.worker.corpus import CorpusMixin
-from arkindex_worker.worker.dataset import DatasetMixin, DatasetState
+from arkindex_worker.worker.dataset import (
+    DatasetMixin,
+    DatasetState,
+    MissingDatasetArchive,
+)
 from arkindex_worker.worker.element import ElementMixin
 from arkindex_worker.worker.entity import EntityMixin
 from arkindex_worker.worker.image import ImageMixin
 from arkindex_worker.worker.metadata import MetaDataMixin, MetaType  # noqa: F401
+from arkindex_worker.worker.process import ActivityState, ProcessMode
 from arkindex_worker.worker.task import TaskMixin
 from arkindex_worker.worker.transcription import TranscriptionMixin
 from arkindex_worker.worker.version import WorkerVersionMixin
 
 
-class ActivityState(Enum):
-    """
-    Processing state of an element.
-    """
-
-    Queued = "queued"
-    """
-    The element has not yet been processed by a worker.
-    """
-
-    Started = "started"
-    """
-    The element is being processed by a worker.
-    """
-
-    Processed = "processed"
-    """
-    The element has been successfully processed by a worker.
-    """
-
-    Error = "error"
-    """
-    An error occurred while processing this element.
-    """
-
-
 class ElementsWorker(
+    ElementMixin,
+    DatasetMixin,
     BaseWorker,
     ClassificationMixin,
     CorpusMixin,
-    ElementMixin,
     TranscriptionMixin,
     WorkerVersionMixin,
     EntityMixin,
@@ -96,22 +74,7 @@ class ElementsWorker(
 
         self._worker_version_cache = {}
 
-    def add_arguments(self):
-        """Define specific ``argparse`` arguments for this worker"""
-        self.parser.add_argument(
-            "--elements-list",
-            help="JSON elements list to use",
-            type=open,
-            default=os.environ.get("TASK_ELEMENTS"),
-        )
-        self.parser.add_argument(
-            "--element",
-            type=str,
-            nargs="+",
-            help="One or more Arkindex element ID",
-        )
-
-    def get_elements(self) -> Iterable[CachedElement] | list[str]:
+    def get_elements(self) -> Iterable[CachedElement] | list[str] | list[Element]:
         """
         List the elements to be processed, either from the CLI arguments or
         the cache database when enabled.
@@ -143,15 +106,20 @@ class ElementsWorker(
         )
         if self.use_cache and cache_query.exists():
             return cache_query
-        # Process elements from JSON file
         elif self.args.elements_list:
+            # Process elements from JSON file
             data = json.load(self.args.elements_list)
             assert isinstance(data, list), "Elements list must be a list"
             assert len(data), "No elements in elements list"
             out += list(filter(None, [element.get("id") for element in data]))
-        # Add any extra element from CLI
         elif self.args.element:
+            # Add any extra element from CLI
             out += self.args.element
+        elif self.process_mode == ProcessMode.Dataset or self.args.set:
+            # Elements from datasets
+            return list(
+                chain.from_iterable(map(self.list_set_elements, self.list_sets()))
+            )
 
         invalid_element_ids = list(filter(invalid_element_id, out))
         assert (
@@ -166,7 +134,8 @@ class ElementsWorker(
         Whether or not WorkerActivity support has been enabled on the DataImport
         used to run this worker.
         """
-        if self.is_read_only:
+        if self.is_read_only or self.process_mode == ProcessMode.Dataset:
+            # Worker activities are also disabled when running an ElementsWorker in a Dataset process.
             return False
         assert (
             self.process_information
@@ -200,7 +169,7 @@ class ElementsWorker(
         for i, item in enumerate(elements, start=1):
             element = None
             try:
-                if self.use_cache:
+                if isinstance(item, CachedElement | Element):
                     # Just use the result of get_elements as the element
                     element = item
                 else:
@@ -316,29 +285,7 @@ class ElementsWorker(
         return True
 
 
-def check_dataset_set(value: str) -> tuple[uuid.UUID, str]:
-    values = value.split(":")
-    if len(values) != 2:
-        raise ArgumentTypeError(
-            f"'{value}' is not in the correct format `<dataset_id>:<set_name>`"
-        )
-
-    dataset_id, set_name = values
-    try:
-        dataset_id = uuid.UUID(dataset_id)
-        return (dataset_id, set_name)
-    except (TypeError, ValueError) as e:
-        raise ArgumentTypeError(f"'{dataset_id}' should be a valid UUID") from e
-
-
-class MissingDatasetArchive(Exception):
-    """
-    Exception raised when the compressed archive associated to
-    a dataset isn't found in its task artifacts.
-    """
-
-
-class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
+class DatasetWorker(DatasetMixin, BaseWorker, TaskMixin):
     """
     Base class for ML workers that operate on Arkindex dataset sets.
 
@@ -361,19 +308,6 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
         # Set as an instance variable as dataset workers might use it to easily extract its content
         self.downloaded_dataset_artifact: Path | None = None
 
-    def add_arguments(self):
-        """Define specific ``argparse`` arguments for this worker"""
-        self.parser.add_argument(
-            "--set",
-            type=check_dataset_set,
-            nargs="+",
-            help="""
-                One or more Arkindex dataset sets, format is <dataset_uuid>:<set_name>
-                (e.g.: "12341234-1234-1234-1234-123412341234:train")
-            """,
-            default=[],
-        )
-
     def cleanup_downloaded_artifact(self) -> None:
         """
         Cleanup the downloaded dataset artifact if any
@@ -421,30 +355,10 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
         :param set: The set to process.
         """
 
-    def list_sets(self) -> Iterator[Set]:
-        """
-        List the sets to be processed, either from the CLI arguments or using the
-        [list_process_sets][arkindex_worker.worker.dataset.DatasetMixin.list_process_sets] method.
-
-        :returns: An iterator of ``Set`` objects.
-        """
-        if not self.is_read_only:
-            yield from self.list_process_sets()
-
-        datasets: dict[uuid.UUID, Dataset] = {}
-        for dataset_id, set_name in self.args.set:
-            # Retrieving dataset information is not already cached
-            if dataset_id not in datasets:
-                datasets[dataset_id] = Dataset(
-                    **self.api_client.request("RetrieveDataset", id=dataset_id)
-                )
-
-            yield Set(name=set_name, dataset=datasets[dataset_id])
-
     def run(self):
         """
         Implements an Arkindex worker that goes through each dataset set returned by
-        [list_sets][arkindex_worker.worker.DatasetWorker.list_sets].
+        [list_sets][arkindex_worker.worker.dataset.DatasetMixin.list_sets].
 
         It calls [process_set][arkindex_worker.worker.DatasetWorker.process_set],
         catching exceptions.

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/arkindex_worker/worker/base.py

@@ -24,6 +24,7 @@ from arkindex_worker.cache import (
     merge_parents_cache,
 )
 from arkindex_worker.utils import close_delete_file, extract_tar_zst_archive
+from arkindex_worker.worker.process import ProcessMode
 from teklia_toolbox.requests import get_arkindex_client
 
 
@@ -156,6 +157,13 @@ class BaseWorker:
             raise Exception("Missing ARKINDEX_CORPUS_ID environment variable")
         return self._corpus_id
 
+    @property
+    def process_mode(self) -> ProcessMode | None:
+        """Mode of the process being run. Returns None when read-only."""
+        if self.is_read_only:
+            return
+        return ProcessMode(self.process_information["mode"])
+
     @property
     def is_read_only(self) -> bool:
         """

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/arkindex_worker/worker/dataset.py

@@ -2,6 +2,8 @@
 BaseWorker methods for datasets.
 """
 
+import uuid
+from argparse import ArgumentTypeError
 from collections.abc import Iterator
 from enum import Enum
 
@@ -36,7 +38,55 @@ class DatasetState(Enum):
     """
 
 
+class MissingDatasetArchive(Exception):
+    """
+    Exception raised when the compressed archive associated to
+    a dataset isn't found in its task artifacts.
+    """
+
+
+def check_dataset_set(value: str) -> tuple[uuid.UUID, str]:
+    """The `--set` argument should have the following format:
+    <dataset_id>:<set_name>
+
+    Args:
+        value (str): Provided argument.
+
+    Raises:
+        ArgumentTypeError: When the value is invalid.
+
+    Returns:
+        tuple[uuid.UUID, str]: The ID of the dataset parsed as UUID and the name of the set.
+    """
+    values = value.split(":")
+    if len(values) != 2:
+        raise ArgumentTypeError(
+            f"'{value}' is not in the correct format `<dataset_id>:<set_name>`"
+        )
+
+    dataset_id, set_name = values
+    try:
+        dataset_id = uuid.UUID(dataset_id)
+        return (dataset_id, set_name)
+    except (TypeError, ValueError) as e:
+        raise ArgumentTypeError(f"'{dataset_id}' should be a valid UUID") from e
+
+
 class DatasetMixin:
+    def add_arguments(self) -> None:
+        """Define specific ``argparse`` arguments for the worker using this mixin"""
+        self.parser.add_argument(
+            "--set",
+            type=check_dataset_set,
+            nargs="+",
+            help="""
+                One or more Arkindex dataset sets, format is <dataset_uuid>:<set_name>
+                (e.g.: "12341234-1234-1234-1234-123412341234:train")
+            """,
+            default=[],
+        )
+        super().add_arguments()
+
     def list_process_sets(self) -> Iterator[Set]:
         """
         List dataset sets associated to the worker's process. This helper is not available in developer mode.
@@ -73,6 +123,26 @@ class DatasetMixin:
 
         return map(lambda result: Element(**result["element"]), results)
 
+    def list_sets(self) -> Iterator[Set]:
+        """
+        List the sets to be processed, either from the CLI arguments or using the
+        [list_process_sets][arkindex_worker.worker.dataset.DatasetMixin.list_process_sets] method.
+
+        :returns: An iterator of ``Set`` objects.
+        """
+        if not self.is_read_only:
+            yield from self.list_process_sets()
+
+        datasets: dict[uuid.UUID, Dataset] = {}
+        for dataset_id, set_name in self.args.set:
+            # Retrieving dataset information if not already cached
+            if dataset_id not in datasets:
+                datasets[dataset_id] = Dataset(
+                    **self.api_client.request("RetrieveDataset", id=dataset_id)
+                )
+
+            yield Set(name=set_name, dataset=datasets[dataset_id])
+
     @unsupported_cache
     def update_dataset_state(self, dataset: Dataset, state: DatasetState) -> Dataset:
         """

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/arkindex_worker/worker/element.py

@@ -2,6 +2,7 @@
 ElementsWorker methods for elements and element types.
 """
 
+import os
 from collections.abc import Iterable
 from operator import attrgetter
 from typing import NamedTuple
@@ -38,6 +39,22 @@ class MissingTypeError(Exception):
 
 
 class ElementMixin:
+    def add_arguments(self):
+        """Define specific ``argparse`` arguments for the worker using this mixin"""
+        self.parser.add_argument(
+            "--elements-list",
+            help="JSON elements list to use",
+            type=open,
+            default=os.environ.get("TASK_ELEMENTS"),
+        )
+        self.parser.add_argument(
+            "--element",
+            type=str,
+            nargs="+",
+            help="One or more Arkindex element ID",
+        )
+        super().add_arguments()
+
     def list_corpus_types(self):
         """
         Loads available element types in corpus.

arkindex_base_worker-0.4.0rc2/arkindex_worker/worker/process.py

@@ -0,0 +1,63 @@
+from enum import Enum
+
+
+class ActivityState(Enum):
+    """
+    Processing state of an element.
+    """
+
+    Queued = "queued"
+    """
+    The element has not yet been processed by a worker.
+    """
+
+    Started = "started"
+    """
+    The element is being processed by a worker.
+    """
+
+    Processed = "processed"
+    """
+    The element has been successfully processed by a worker.
+    """
+
+    Error = "error"
+    """
+    An error occurred while processing this element.
+    """
+
+
+class ProcessMode(Enum):
+    """
+    Mode of the process of the worker.
+    """
+
+    Files = "files"
+    """
+    Processes of files (images, PDFs, IIIF, ...) imports.
+    """
+
+    Workers = "workers"
+    """
+    Processes of worker executions.
+    """
+
+    Template = "template"
+    """
+    Process templates.
+    """
+
+    S3 = "s3"
+    """
+    Processes of imports from an S3-compatible storage.
+    """
+
+    Local = "local"
+    """
+    Local processes.
+    """
+
+    Dataset = "dataset"
+    """
+    Dataset processes.
+    """

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/pyproject.toml

@@ -4,11 +4,11 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "arkindex-base-worker"
-version = "0.4.0rc1"
+version = "0.4.0rc2"
 description = "Base Worker to easily build Arkindex ML workflows"
 license = { file = "LICENSE" }
 dependencies = [
-    "humanize==4.9.0",
+    "humanize==4.10.0",
     "peewee~=3.17",
     "Pillow==10.4.0",
     "python-gnupg==0.5.2",
@@ -44,7 +44,7 @@ Authors = "https://teklia.com"
 [project.optional-dependencies]
 docs = [
     "black==24.4.2",
-    "mkdocs-material==9.5.31",
+    "mkdocs-material==9.5.33",
     "mkdocstrings-python==1.10.8",
 ]
 tests = [

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/tests/__init__.py

@@ -5,4 +5,4 @@ FIXTURES_DIR = BASE_DIR / "data"
 SAMPLES_DIR = BASE_DIR / "samples"
 
 CORPUS_ID = "11111111-1111-1111-1111-111111111111"
-PROCESS_ID = "cafecafe-cafe-cafe-cafe-cafecafecafe"
+PROCESS_ID = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/tests/conftest.py

@@ -23,10 +23,15 @@ from arkindex_worker.cache import (
     init_cache_db,
 )
 from arkindex_worker.models import Artifact, Dataset, Set
-from arkindex_worker.worker import BaseWorker, DatasetWorker, ElementsWorker
+from arkindex_worker.worker import (
+    BaseWorker,
+    DatasetWorker,
+    ElementsWorker,
+    ProcessMode,
+)
 from arkindex_worker.worker.dataset import DatasetState
 from arkindex_worker.worker.transcription import TextOrientation
-from tests import CORPUS_ID, PROCESS_ID, SAMPLES_DIR
+from tests import CORPUS_ID, SAMPLES_DIR
 
 __yaml_cache = {}
 
@@ -601,7 +606,9 @@ def mock_dataset_worker(monkeypatch, mocker, _mock_worker_run_api):
 
     dataset_worker = DatasetWorker()
     dataset_worker.configure()
-    dataset_worker.process_information = {"id": PROCESS_ID}
+
+    # Update process mode
+    dataset_worker.process_information["mode"] = ProcessMode.Dataset
 
     assert not dataset_worker.is_read_only
 

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/tests/test_dataset_worker.py

@@ -6,8 +6,11 @@ import pytest
 from apistar.exceptions import ErrorResponse
 
 from arkindex_worker.models import Dataset, Set
-from arkindex_worker.worker import MissingDatasetArchive, check_dataset_set
-from arkindex_worker.worker.dataset import DatasetState
+from arkindex_worker.worker.dataset import (
+    DatasetState,
+    MissingDatasetArchive,
+    check_dataset_set,
+)
 from tests import FIXTURES_DIR, PROCESS_ID
 from tests.test_elements_worker import BASE_API_CALLS
 

{arkindex_base_worker-0.4.0rc1 → arkindex_base_worker-0.4.0rc2}/tests/test_elements_worker/test_elements.py

@@ -17,7 +17,9 @@ from arkindex_worker.cache import (
 from arkindex_worker.models import Element
 from arkindex_worker.utils import DEFAULT_BATCH_SIZE
 from arkindex_worker.worker import ElementsWorker
+from arkindex_worker.worker.dataset import DatasetState
 from arkindex_worker.worker.element import MissingTypeError
+from arkindex_worker.worker.process import ProcessMode
 from tests import CORPUS_ID
 
 from . import BASE_API_CALLS
@@ -208,10 +210,12 @@ def test_get_elements_element_arg_not_uuid(mocker, mock_elements_worker):
         "arkindex_worker.worker.base.argparse.ArgumentParser.parse_args",
         return_value=Namespace(
             element=["volumeid", "pageid"],
+            config={},
             verbose=False,
             elements_list=None,
             database=None,
-            dev=False,
+            dev=True,
+            set=[],
         ),
     )
 
@@ -232,10 +236,12 @@ def test_get_elements_element_arg(mocker, mock_elements_worker):
                 "11111111-1111-1111-1111-111111111111",
                 "22222222-2222-2222-2222-222222222222",
             ],
+            config={},
             verbose=False,
             elements_list=None,
             database=None,
-            dev=False,
+            dev=True,
+            set=[],
         ),
     )
 
@@ -250,6 +256,264 @@ def test_get_elements_element_arg(mocker, mock_elements_worker):
     ]
 
 
+def test_get_elements_dataset_set_arg(responses, mocker, mock_elements_worker):
+    mocker.patch(
+        "arkindex_worker.worker.base.argparse.ArgumentParser.parse_args",
+        return_value=Namespace(
+            element=[],
+            config={},
+            verbose=False,
+            elements_list=None,
+            database=None,
+            dev=True,
+            set=[(UUID("11111111-1111-1111-1111-111111111111"), "train")],
+        ),
+    )
+
+    # Mock RetrieveDataset call
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/datasets/11111111-1111-1111-1111-111111111111/",
+        status=200,
+        json={
+            "id": "11111111-1111-1111-1111-111111111111",
+            "name": "My dataset",
+            "description": "A dataset about cats.",
+            "sets": ["train", "dev", "test"],
+            "state": DatasetState.Complete.value,
+        },
+        content_type="application/json",
+    )
+
+    # Mock ListSetElements call
+    element = {
+        "id": "22222222-2222-2222-2222-222222222222",
+        "type": "page",
+        "name": "1",
+        "corpus": {
+            "id": "11111111-1111-1111-1111-111111111111",
+        },
+        "thumbnail_url": "http://example.com",
+        "zone": {
+            "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "polygon": [[0, 0], [0, 0], [0, 0]],
+            "image": {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "path": "string",
+                "width": 0,
+                "height": 0,
+                "url": "http://example.com",
+                "s3_url": "string",
+                "status": "checked",
+                "server": {
+                    "display_name": "string",
+                    "url": "http://example.com",
+                    "max_width": 2147483647,
+                    "max_height": 2147483647,
+                },
+            },
+            "url": "http://example.com",
+        },
+        "rotation_angle": 0,
+        "mirrored": False,
+        "created": "2019-08-24T14:15:22Z",
+        "classes": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "ml_class": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "name": "string",
+                },
+                "state": "pending",
+                "confidence": 0,
+                "high_confidence": True,
+                "worker_run": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "summary": "string",
+                },
+            }
+        ],
+        "metadata": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "type": "text",
+                "name": "string",
+                "value": "string",
+                "dates": [{"type": "exact", "year": 0, "month": 1, "day": 1}],
+            }
+        ],
+        "transcriptions": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "text": "string",
+                "confidence": 0,
+                "orientation": "horizontal-lr",
+                "worker_run": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "summary": "string",
+                },
+            }
+        ],
+        "has_children": True,
+        "worker_run": {
+            "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "summary": "string",
+        },
+        "confidence": 1,
+    }
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/datasets/11111111-1111-1111-1111-111111111111/elements/?set=train&with_count=true",
+        status=200,
+        json={
+            "next": None,
+            "previous": None,
+            "results": [
+                {
+                    "set": "train",
+                    "element": element,
+                }
+            ],
+            "count": 1,
+        },
+        content_type="application/json",
+    )
+
+    worker = ElementsWorker()
+    worker.configure()
+
+    elt_list = worker.get_elements()
+
+    assert elt_list == [
+        Element(**element),
+    ]
+
+
+def test_get_elements_dataset_set_api(responses, mocker, mock_elements_worker):
+    # Mock ListProcessSets call
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/process/aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff/sets/",
+        status=200,
+        json={
+            "next": None,
+            "previous": None,
+            "results": [
+                {
+                    "id": "33333333-3333-3333-3333-333333333333",
+                    "dataset": {"id": "11111111-1111-1111-1111-111111111111"},
+                    "set_name": "train",
+                }
+            ],
+            "count": 1,
+        },
+        content_type="application/json",
+    )
+
+    # Mock ListSetElements call
+    element = {
+        "id": "22222222-2222-2222-2222-222222222222",
+        "type": "page",
+        "name": "1",
+        "corpus": {
+            "id": "11111111-1111-1111-1111-111111111111",
+        },
+        "thumbnail_url": "http://example.com",
+        "zone": {
+            "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "polygon": [[0, 0], [0, 0], [0, 0]],
+            "image": {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "path": "string",
+                "width": 0,
+                "height": 0,
+                "url": "http://example.com",
+                "s3_url": "string",
+                "status": "checked",
+                "server": {
+                    "display_name": "string",
+                    "url": "http://example.com",
+                    "max_width": 2147483647,
+                    "max_height": 2147483647,
+                },
+            },
+            "url": "http://example.com",
+        },
+        "rotation_angle": 0,
+        "mirrored": False,
+        "created": "2019-08-24T14:15:22Z",
+        "classes": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "ml_class": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "name": "string",
+                },
+                "state": "pending",
+                "confidence": 0,
+                "high_confidence": True,
+                "worker_run": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "summary": "string",
+                },
+            }
+        ],
+        "metadata": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "type": "text",
+                "name": "string",
+                "value": "string",
+                "dates": [{"type": "exact", "year": 0, "month": 1, "day": 1}],
+            }
+        ],
+        "transcriptions": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "text": "string",
+                "confidence": 0,
+                "orientation": "horizontal-lr",
+                "worker_run": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "summary": "string",
+                },
+            }
+        ],
+        "has_children": True,
+        "worker_run": {
+            "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "summary": "string",
+        },
+        "confidence": 1,
+    }
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/datasets/11111111-1111-1111-1111-111111111111/elements/?set=train&with_count=true",
+        status=200,
+        json={
+            "next": None,
+            "previous": None,
+            "results": [
+                {
+                    "set": "train",
+                    "element": element,
+                }
+            ],
+            "count": 1,
+        },
+        content_type="application/json",
+    )
+
+    # Update ProcessMode to Dataset
+    mock_elements_worker.process_information["mode"] = ProcessMode.Dataset
+
+    elt_list = mock_elements_worker.get_elements()
+
+    assert elt_list == [
+        Element(**element),
+    ]
+
+
 def test_get_elements_both_args_error(mocker, mock_elements_worker, tmp_path):
     elements_path = tmp_path / "elements.json"
     elements_path.write_text(
@@ -270,6 +534,7 @@ def test_get_elements_both_args_error(mocker, mock_elements_worker, tmp_path):
             elements_list=elements_path.open(),
             database=None,
             dev=False,
+            set=[],
         ),
     )
 
@@ -295,6 +560,7 @@ def test_database_arg(mocker, mock_elements_worker, tmp_path):
             elements_list=None,
             database=database_path,
             dev=False,
+            set=[],
         ),
     )
 
@@ -319,6 +585,7 @@ def test_database_arg_cache_missing_version_table(
             elements_list=None,
             database=database_path,
             dev=False,
+            set=[],
         ),
     )