Rework the worker due to `Dataset` API changes

Closes #14 (closed), #16 (closed)

tests/test_worker.py

@@ -4,6 +4,7 @@ import json
 from argparse import Namespace
 from uuid import UUID, uuid4
 
+from arkindex_export.models import Element
 from arkindex_worker.cache import (
     CachedClassification,
     CachedDataset,
@@ -14,7 +15,6 @@ from arkindex_worker.cache import (
     CachedTranscription,
     CachedTranscriptionEntity,
 )
-from worker_generic_training_dataset.db import retrieve_element
 from worker_generic_training_dataset.worker import DatasetExtractor
 
 
@@ -44,8 +44,8 @@ def test_process_split(tmp_path, downloaded_images):
     worker.process_split(
         "train",
         [
-            retrieve_element(first_page_id),
-            retrieve_element(second_page_id),
+            Element.get_by_id(first_page_id),
+            Element.get_by_id(second_page_id),
         ],
     )