Merge pull request #441 from RelevanceAI/development

v0.32.1

Author: jtwinrelevanceai

ai_transform/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.32.0"
+__version__ = "0.32.1"
 
 from ai_transform.timer import Timer
 

ai_transform/api/api.py

@@ -162,9 +162,10 @@ def _list_datasets(self):
     def _create_dataset(
         self, dataset_id: str, schema: Optional[Schema] = None, upsert: bool = True, expire: bool = False
     ) -> Any:
-        response = self.post(
-            suffix=f"/datasets/create", json=dict(id=dataset_id, schema=schema, upsert=upsert, expire=expire)
-        )
+        obj = dict(id=dataset_id, upsert=upsert, expire=expire)
+        if schema:
+            obj["schema"] = schema
+        response = self.post(suffix=f"/datasets/create", json=obj)
         return get_response(response)
 
     def _delete_dataset(self, dataset_id: str) -> Any:

ai_transform/dataset/field.py

@@ -107,11 +107,10 @@ def contains(self, other: str) -> Filter:
 
     def exists(self) -> Filter:
         if "_chunk_" in self._field:
-            count = self._field.count(".")
-            if count:
-                parent_field = self._field.split(".")[0]
-            else:
+            if self._field.endswith("_chunk_"):
                 parent_field = self._field
+            else:
+                parent_field = self._field.split(".")[0]
 
             return [{"chunk": {"path": parent_field, "filters": [{"fieldExists": {"field": self._field}}]}}]
         return [{"field": self._field, "filter_type": "exists", "condition": "==", "condition_value": " "}]

ai_transform/engine/abstract_engine.py

@@ -63,7 +63,7 @@ def __init__(
                     chunk_index = field.index("_chunk_") + len("_chunk_")
                     chunk_field = field[:chunk_index]
                     fields_to_add += [chunk_field]
-            select_fields += fields_to_add
+            select_fields = select_fields + fields_to_add
             select_fields = list(set(select_fields))
         else:
             select_fields = []
@@ -109,8 +109,8 @@ def __init__(
         self._refresh = refresh
         self._after_id = after_id
 
-        filters += self._get_refresh_filter()
-        filters += self._get_workflow_filter()
+        filters = filters + self._get_refresh_filter()
+        filters = filters + self._get_workflow_filter()
 
         self._filters = filters
 

ai_transform/engine/dense_output_engine.py

@@ -21,6 +21,8 @@
 
 
 class DenseOutputEngine(AbstractEngine):
+    operator: DenseOperator
+
     def __init__(
         self,
         dataset: Dataset = None,
@@ -72,6 +74,7 @@ def apply(self) -> None:
         for mega_batch in self.api_progress(iterator):
             for mini_batch in AbstractEngine.chunk_documents(self._transform_chunksize, mega_batch):
                 document_mapping = self._operate(mini_batch)
+
                 for dataset_id, documents in document_mapping.items():
                     output_dataset_ids.append(dataset_id)
                     dataset = Dataset.from_details(dataset_id, self.token)
@@ -81,7 +84,14 @@ def apply(self) -> None:
         self.operator.post_hooks(self._dataset)
 
         output_datasets = self.datasets_from_ids(output_dataset_ids)
-        self.operator.store_dataset_relationship(self.dataset, output_datasets)
+        self.store_dataset_relationship(output_datasets)
 
     def datasets_from_ids(self, dataset_ids: Sequence[str]) -> Sequence[Dataset]:
         return [Dataset.from_details(dataset_id, self.token) for dataset_id in dataset_ids]
+
+    def store_dataset_relationship(self, output_datasets: Sequence[Dataset]):
+        self.dataset.update_metadata(
+            {"_child_datasets_": [output_dataset.dataset_id for output_dataset in output_datasets]}
+        )
+        for output_dataset in output_datasets:
+            output_dataset.update_metadata({"_parent_dataset_": self.dataset.dataset_id})

ai_transform/operator/abstract_operator.py

@@ -75,7 +75,6 @@ def __init__(
         output_fields: Optional[Union[Dict[str, str], List[str]]] = None,
         enable_postprocess: Optional[bool] = True,
     ):
-
         if input_fields is not None and output_fields is not None:
             if any(input_field in output_fields for input_field in input_fields):
                 detected_fields = [input_field for input_field in input_fields if input_field in output_fields]
@@ -171,6 +170,12 @@ def transform_for_playground(
         from ai_transform.api.client import Client
 
         output = self.transform(documents=documents)
+        if hasattr(documents, "to_json"):
+            output = output.to_json()
+        else:
+            for index in range(len(output)):
+                if hasattr(output[index], "to_json"):
+                    output[index] = output[index].to_json()
         client = Client(authorization_token)
         return client.api._set_workflow_status(
             job_id=job_id,
@@ -180,7 +185,7 @@ def transform_for_playground(
             status=status,
             send_email=send_email,
             worker_number=worker_number,
-            output=output,
+            output={"output": output},
         )
 
     def pre_hooks(self, dataset: Dataset):

ai_transform/operator/dense_operator.py

@@ -5,7 +5,6 @@
 from typing import Dict, Sequence
 
 from ai_transform.operator.abstract_operator import AbstractOperator
-from ai_transform.dataset.dataset import Dataset
 from ai_transform.utils.document import Document
 from ai_transform.utils.document_list import DocumentList
 
@@ -24,18 +23,13 @@
 class DenseOperator(AbstractOperator):
     def __call__(self, old_documents: DocumentList) -> DenseOperatorOutput:
         datum = self.transform(old_documents)
-        assert isinstance(datum, dict), BAD_OPERATOR_MESSAGE
+        if not isinstance(datum, dict):
+            raise ValueError(BAD_OPERATOR_MESSAGE)
         for _, documents in datum.items():
-            assert isinstance(documents, Sequence)
+            if not isinstance(documents, Sequence):
+                raise ValueError(BAD_OPERATOR_MESSAGE)
         return datum
 
     @abstractmethod
     def transform(self, documents: DocumentList) -> DenseOperatorOutput:
         raise NotImplementedError
-
-    def store_dataset_relationship(self, input_dataset: Dataset, output_datasets: Sequence[Dataset]):
-        input_dataset.update_metadata(
-            {"_child_datasets_": [output_dataset.dataset_id for output_dataset in output_datasets]}
-        )
-        for output_dataset in output_datasets:
-            output_dataset.update_metadata({"_parent_dataset_": input_dataset.dataset_id})

ai_transform/utils/document.py

@@ -141,17 +141,7 @@ def list_chunks(self):
         return [k for k in self.keys() if k.endswith("_chunk_")]
 
     def get_chunk(self, chunk_field: str, field: str = None, default: str = None):
-        """
-        Returns a list of values.
-        """
-        # provide a recursive implementation for getting chunks
-        from ai_transform.utils.document_list import DocumentList
-
-        document_list = DocumentList(self.get(chunk_field, default=default))
-        # Get the field across chunks
-        if field is None:
-            return document_list
-        return [d.get(field, default=default) for d in document_list.data]
+        return [document.get(field, default) for document in self.get(chunk_field, default=default)]
 
     def _create_chunk_documents(self, field: str, values: list, generate_id: bool = False):
         """

ai_transform/utils/example_documents.py

@@ -52,7 +52,6 @@ def create_id():
 
 
 def generate_random_string(string_length: int = 5) -> str:
-
     """Generate a random string of letters and numbers"""
     return "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(string_length))
 
@@ -92,7 +91,16 @@ def vector_document(vector_length: int) -> Document:
 
 
 def mock_documents(n: int = 100, vector_length: int = 5) -> DocumentList:
-    return DocumentList([vector_document(vector_length) for _ in range(n)])
+    documents = [vector_document(vector_length) for _ in range(n)]
+    return DocumentList(documents)
+
+
+def incomplete_documents(n: int = 100, vector_length: int = 5) -> DocumentList:
+    documents = [vector_document(vector_length).data for _ in range(n)]
+    for document in documents:
+        for key in random.sample(document.keys(), 3):
+            document.pop(key)
+    return DocumentList(documents)
 
 
 def static_documents(n: int = 100) -> DocumentList:

examples/workflows/clustering_example.py

@@ -71,7 +71,7 @@ def execute(token: str, logger: Callable, worker_number: int = 0, *args, **kwarg
     engine = InMemoryEngine(
         dataset=dataset,
         operator=operator,
-        chunksize=16,
+        chunksize=8,
         select_fields=[vector_field],
         filters=filters,
         worker_number=worker_number,