Allow using python functions instead of operators (e.g in pre-processing pipeline)

Signed-off-by: elronbandel <[email protected]>

Author: elronbandel

docs/catalog.py

@@ -1,6 +1,7 @@
 import json
 import os
 import re
+import types
 from collections import defaultdict
 from functools import lru_cache
 from pathlib import Path
@@ -110,6 +111,7 @@ def all_subtypes_of_artifact(artifact):
         or isinstance(artifact, bool)
         or isinstance(artifact, int)
         or isinstance(artifact, float)
+        or isinstance(artifact, types.FunctionType)
     ):
         return []
     if isinstance(artifact, list):

docs/docs/adding_dataset.rst

@@ -5,7 +5,7 @@
    To use this tutorial, you need to :ref:`install Unitxt <install_unitxt>`.
 
 =================
-Datasets 
+Datasets
 =================
 
 This guide will assist you in adding or using your new dataset in Unitxt.
@@ -105,6 +105,27 @@ Most data can be normalized to the task schema using built-in operators, ensurin
 
 For custom operators, refer to the :ref:`Operators Tutorial <adding_operator>`.
 
+.. tip::
+
+    If you cannot find operators fit to your needs simply use instance function operator:
+
+    .. code-block:: python
+
+        def my_function(instance, stream_name=None):
+            instance["x"] += 42
+            return instance
+
+    Or stream function operator:
+
+    .. code-block:: python
+
+        def my_other_function(stream, stream_name=None):
+            for instance in stream:
+                instance["x"] += 42
+                yield instance
+
+    Both functions can be plugged in every place in unitxt requires operators, e.g pre-processing pipeline.
+
 The Template
 ----------------
 

docs/docs/adding_operator.rst

@@ -5,11 +5,33 @@
    To use this tutorial, you need to :ref:`install unitxt <install_unitxt>`.
 
 =====================================
-Operators 
+Operators
 =====================================
 
 Operators are specialized functions designed to process data.
 
+.. tip::
+
+    If you cannot find operators fit to your needs simply use instance function operator:
+
+    .. code-block:: python
+
+        def my_function(instance, stream_name=None):
+            instance["x"] += 42
+            return instance
+
+    Or stream function operator:
+
+    .. code-block:: python
+
+        def my_other_function(stream, stream_name=None):
+            for instance in stream:
+                instance["x"] += 42
+                yield instance
+
+    Both functions can be plugged in every place in unitxt requires operators, e.g pre-processing pipeline.
+
+
 They are used in the TaskCard for preparing data for specific tasks and by Post Processors
 to process the textual output of the model to the expect input of the metrics.
 
@@ -18,11 +40,11 @@ There are several types of operators.
 1. Field Operators - Operators that modify individual fields of the instances in the input streams.  Example of such operators are operators that
 cast field values, uppercase string fields, or translate text between languages.
 
-2. Instance Operators - Operators that modify individual instances in the input streams. For example, operators that add or remove fields.
+1. Instance Operators - Operators that modify individual instances in the input streams. For example, operators that add or remove fields.
 
-3. Stream Operators - Operators that perform operations on full streams. For example, operators that remove instances based on some condition.
+2. Stream Operators - Operators that perform operations on full streams. For example, operators that remove instances based on some condition.
 
-4. MultiStream Operators - Operator that perform operations on multiple streams.  For example, operators that repartition the instances between train and test splits.
+3. MultiStream Operators - Operator that perform operations on multiple streams.  For example, operators that repartition the instances between train and test splits.
 
 Unitxt comes with a large collection of built in operators - that were design to cover most common requirements of dataset processing.
 

prepare/cards/xlam_function_calling.py

@@ -3,14 +3,26 @@
 from unitxt.loaders import LoadHF
 from unitxt.operators import (
     Copy,
-    ExecuteExpression,
     Move,
     Set,
 )
 from unitxt.splitters import RenameSplits
 from unitxt.struct_data_operators import LoadJson
 from unitxt.test_utils.card import test_card
 
+
+def extract_required_parameters(instance, stream_name=None):
+    result = []
+    for tool in instance["tools"]:
+        required_params = []
+        for param_name, param_info in tool["parameters"]["properties"].items():
+            if "optional" not in param_info["type"]:
+                required_params.append(param_name)
+        result.append(required_params)
+    instance["required"] = result
+    return instance
+
+
 card = TaskCard(
     loader=LoadHF(
         path="Salesforce/xlam-function-calling-60k",
@@ -30,10 +42,7 @@
             to_field="tools/*/parameters/properties",
             set_every_value=True,
         ),
-        ExecuteExpression(
-            to_field="required",
-            expression="[[p for p, c in tool['parameters']['properties'].items() if 'optional' not in c['type']] for tool in tools]",
-        ),
+        extract_required_parameters,
         Copy(
             field="required",
             to_field="tools/*/parameters/required",

src/unitxt/api.py

@@ -27,6 +27,7 @@
 from .settings_utils import get_constants, get_settings
 from .standard import DatasetRecipe
 from .task import Task
+from .utils import json_dump
 
 logger = get_logger()
 constants = get_constants()
@@ -180,6 +181,20 @@ class MyClass:
     return obj_str
 
 
+def _remove_id_keys(obj):
+    if isinstance(obj, dict):
+        return {k: _remove_id_keys(v) for k, v in obj.items() if k != "__id__"}
+    if isinstance(obj, list):
+        return [_remove_id_keys(item) for item in obj]
+    return obj
+
+
+def _artifact_string_repr(artifact):
+    artifact_dict = to_dict(artifact, object_to_str_without_addresses)
+    artifact_dict_without_ids = _remove_id_keys(artifact_dict)
+    return json_dump(artifact_dict_without_ids)
+
+
 def _source_to_dataset(
     source: SourceOperator,
     split=None,
@@ -189,9 +204,7 @@ def _source_to_dataset(
     from .dataset import Dataset as UnitxtDataset
 
     # Generate a unique signature for the source
-    source_signature = json.dumps(
-        to_dict(source, object_to_str_without_addresses), sort_keys=True
-    )
+    source_signature = _artifact_string_repr(source)
     config_name = "recipe-" + short_hex_hash(source_signature)
     # Obtain data stream from the source
     stream = source()

src/unitxt/artifact.py

@@ -4,6 +4,7 @@
 import os
 import pkgutil
 import re
+import types
 import warnings
 from abc import abstractmethod
 from typing import Any, Dict, List, Optional, Tuple, Union, final
@@ -27,6 +28,8 @@
 from .utils import (
     artifacts_json_cache,
     json_dump,
+    json_load,
+    load_json,
     save_to_file,
     shallow_copy,
 )
@@ -119,16 +122,32 @@ def reset(self):
         self.catalogs = []
 
 
+def maybe_recover_function_operator(func):
+    sig = inspect.signature(func)
+    param_names = tuple(sorted(sig.parameters))
+    if param_names == ("stream", "stream_name") or param_names == (
+        "instance",
+        "stream_name",
+    ):
+        from .operators import FunctionOperator
+
+        return FunctionOperator(function=func)
+    return func
+
+
 def maybe_recover_artifacts_structure(obj):
+    if isinstance(obj, types.FunctionType):
+        obj = maybe_recover_function_operator(obj)
+
     if Artifact.is_possible_identifier(obj):
         return verbosed_fetch_artifact(obj)
     if isinstance(obj, dict):
         for key, value in obj.items():
-            obj[key] = maybe_recover_artifact(value)
+            obj[key] = maybe_recover_artifacts_structure(value)
         return obj
     if isinstance(obj, list):
         for i in range(len(obj)):
-            obj[i] = maybe_recover_artifact(obj[i])
+            obj[i] = maybe_recover_artifacts_structure(obj[i])
         return obj
     return obj
 
@@ -237,8 +256,7 @@ def __init_subclass__(cls, **kwargs):
     def is_artifact_file(cls, path):
         if not os.path.exists(path) or not os.path.isfile(path):
             return False
-        with open(path) as f:
-            d = json.load(f)
+        d = load_json(path)
         return cls.is_artifact_dict(d)
 
     @classmethod
@@ -384,14 +402,15 @@ def serialize(self):
         return self.to_json()
 
     def save(self, path):
-        original_args = Artifact.from_dict(self.to_dict()).get_repr_dict()
+        data = self.to_dict()
+        original_args = Artifact.from_dict(data).get_repr_dict()
         current_args = self.get_repr_dict()
         diffs = dict_diff_string(original_args, current_args)
         if diffs:
             raise UnitxtError(
                 f"Cannot save catalog artifacts that have changed since initialization. Detected differences in the following fields:\n{diffs}"
             )
-        save_to_file(path, self.to_json())
+        save_to_file(path, json_dump(data))
 
     def verify_instance(
         self, instance: Dict[str, Any], name: Optional[str] = None
@@ -581,7 +600,7 @@ def fetch_artifact(
 
     # If Json string, first load into dictionary
     if isinstance(artifact_rep, str):
-        artifact_rep = json.loads(artifact_rep)
+        artifact_rep = json_load(artifact_rep)
     # Load from dictionary (fails if not valid dictionary)
     return Artifact.from_dict(artifact_rep), None
 
@@ -657,7 +676,7 @@ def get_artifacts_data_classification(artifact: str) -> Optional[List[str]]:
     )
 
     try:
-        data_classification = json.loads(data_classification)
+        data_classification = json_load(data_classification)
     except json.decoder.JSONDecodeError as e:
         raise RuntimeError(error_msg) from e
 

src/unitxt/catalog.py

@@ -18,6 +18,7 @@
 from .logging_utils import get_logger
 from .settings_utils import get_constants
 from .text_utils import print_dict
+from .utils import json_load
 from .version import version
 
 logger = get_logger()
@@ -228,7 +229,7 @@ def _get_tags_from_file(file_path):
     result = Counter()
 
     with open(file_path) as f:
-        data = json.load(f)
+        data = json_load(f)
         if "__tags__" in data and isinstance(data["__tags__"], dict):
             tags = data["__tags__"]
             for key, value in tags.items():

src/unitxt/catalog/cards/xlam_function_calling_60k.json

@@ -61,9 +61,8 @@
             "set_every_value": true
         },
         {
-            "__type__": "execute_expression",
-            "to_field": "required",
-            "expression": "[[p for p, c in tool['parameters']['properties'].items() if 'optional' not in c['type']] for tool in tools]"
+            "__function__": "extract_required_parameters",
+            "source": "def extract_required_parameters(instance, stream_name=None):\n    result = []\n    for tool in instance[\"tools\"]:\n        required_params = []\n        for param_name, param_info in tool['parameters']['properties'].items():\n            if 'optional' not in param_info['type']:\n                required_params.append(param_name)\n        result.append(required_params)\n    instance[\"required\"] = result\n    return instance\n"
         },
         {
             "__type__": "copy",

src/unitxt/inference.py

@@ -51,7 +51,7 @@
 from .operators import ArtifactFetcherMixin
 from .settings_utils import get_constants, get_settings
 from .type_utils import isoftype
-from .utils import retry_connection_with_exponential_backoff
+from .utils import json_load, retry_connection_with_exponential_backoff
 
 constants = get_constants()
 settings = get_settings()
@@ -391,7 +391,7 @@ def to_tools(self, instance):
         if task_data is None:
             return None
         if isinstance(task_data, str):
-            task_data = json.loads(task_data)
+            task_data = json_load(task_data)
         if "__tools__" in task_data:
             return task_data["__tools__"]
         return None
@@ -814,7 +814,7 @@ def _infer_fn(
                 if "task_data" in instance and "__tools__" in instance["task_data"]:
                     task_data = instance["task_data"]
                     if isinstance(task_data, str):
-                        task_data = json.loads(task_data)
+                        task_data = json_load(task_data)
                     tools.append(task_data["__tools__"])
                 else:
                     tools.append(None)
@@ -2768,7 +2768,7 @@ class WMLInferenceEngineChat(WMLInferenceEngineBase, WMLChatParamsMixin):
     def _extract_queries(instance: Dict[str, Any]) -> Tuple[Optional[str], List]:
         task_data = instance["task_data"]
         if isinstance(task_data, str):
-            task_data = json.loads(task_data)
+            task_data = json_load(task_data)
         question = task_data.get("question")
 
         images = [None]
@@ -2888,7 +2888,7 @@ def to_tools(
             return {"tools": None, "tool_choice": None}
 
         if isinstance(task_data, str):
-            task_data = json.loads(task_data)
+            task_data = json_load(task_data)
         if "__tools__" in task_data:
             tools: List[Dict[str, str]] = task_data["__tools__"]
             tool_choice: Optional[Dict[str, str]] = task_data.get("__tool_choice__")
@@ -3186,7 +3186,7 @@ def _infer(
             task_data = instance["task_data"]
 
             if isinstance(task_data, str):
-                task_data = json.loads(task_data)
+                task_data = json_load(task_data)
 
             for option in task_data["options"]:
                 requests.append(
@@ -3866,7 +3866,7 @@ def _infer(
         return_meta_data: bool = False,
     ) -> Union[List[str], List[TextGenerationInferenceOutput]]:
         task_data = [
-            json.loads(instance["task_data"]) if "task_data" in instance else {}
+            json_load(instance["task_data"]) if "task_data" in instance else {}
             for instance in dataset
         ]
         predictions = [td[self.prediction_field] for td in task_data]

src/unitxt/llm_as_judge_from_template.py

@@ -12,16 +12,15 @@
 from .settings_utils import get_settings
 from .system_prompts import EmptySystemPrompt, SystemPrompt
 from .templates import Template
+from .utils import json_load
 
 settings = get_settings()
 
 
 def get_task_data_dict(task_data):
-    import json
-
     # seems like the task data sometimes comes as a string, not a dict
     # this fixes it
-    return json.loads(task_data) if isinstance(task_data, str) else task_data
+    return json_load(task_data) if isinstance(task_data, str) else task_data
 
 
 class LLMAsJudgeBase(BulkInstanceMetric, ArtifactFetcherMixin):