diff --git a/README.rst b/README.rst index df67363..c2db2ca 100644 --- a/README.rst +++ b/README.rst @@ -528,6 +528,19 @@ __ https://docs.scale.com/reference client.create_evaluation_task(TaskType.TextCollection, **payload) +Training tasks (For Scale Rapid projects only) +________________________________________________ + +Training tasks are used to onboard taskers onto your project + +Create Training Task +^^^^^^^^^^^^^^^^^^^^^^ + +Create a training task. + +.. code-block:: python + + client.create_training_task(TaskType, ...task parameters...) Error handling ______________ diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py index 58e25f4..2958c47 100644 --- a/scaleapi/__init__.py +++ b/scaleapi/__init__.py @@ -5,6 +5,7 @@ from scaleapi.exceptions import ScaleInvalidRequest from scaleapi.files import File from scaleapi.projects import Project +from scaleapi.training_tasks import TrainingTask from ._version import __version__ # noqa: F401 from .api import Api @@ -794,10 +795,12 @@ def create_evaluation_task( task_type: TaskType, **kwargs, ) -> EvaluationTask: - """This method can only be used for Self-Serve projects. + """This method can only be used for Rapid projects. Supported Task Types: [ + DocumentTranscription, + SegmentAnnotation, + VideoPlaybackAnnotation, ImageAnnotation, - Categorization, TextCollection, NamedEntityRecognition ] @@ -827,3 +830,38 @@ def create_evaluation_task( evaluation_task_data = self.api.post_request(endpoint, body=kwargs) return EvaluationTask(evaluation_task_data, self) + + def create_training_task( + self, + task_type: TaskType, + **kwargs, + ) -> TrainingTask: + """This method can only be used for Rapid projects. + Supported Task Types: [ + DocumentTranscription, + SegmentAnnotation, + VideoPlaybackAnnotation, + ImageAnnotation, + TextCollection, + NamedEntityRecognition + ] + Parameters may differ based on the given task_type. + + Args: + task_type (TaskType): + Task type to be created + e.g.. `TaskType.ImageAnnotation` + **kwargs: + The same set of parameters are expected with + create_task function and an additional expected_response. + Scale's API documentation. + https://docs.scale.com/reference + + Returns: + TrainingTask: + Returns created training task. + """ + endpoint = f"training_tasks/{task_type.value}" + + training_task_data = self.api.post_request(endpoint, body=kwargs) + return TrainingTask(training_task_data, self) diff --git a/scaleapi/training_tasks.py b/scaleapi/training_tasks.py new file mode 100644 index 0000000..b1f3895 --- /dev/null +++ b/scaleapi/training_tasks.py @@ -0,0 +1,22 @@ +class TrainingTask: + """TrainingTask class, containing TrainingTask information.""" + + def __init__(self, json, client): + self._json = json + self.id = json["id"] + self.initial_response = getattr(json, "initial_response", None) + self.expected_response = json["expected_response"] + self._client = client + + def __hash__(self): + return hash(self.id) + + def __str__(self): + return f"TrainingTask(id={self.id})" + + def __repr__(self): + return f"TrainingTask({self._json})" + + def as_dict(self): + """Returns all attributes as a dictionary""" + return self._json