Add an ability to create a model chat evaluation ontology

Add create_model_chat_evaluation_ontology_from_feature_schemas

Author: Val Brodsky

libs/labelbox/src/labelbox/client.py

@@ -52,6 +52,7 @@
 from labelbox.schema.slice import CatalogSlice, ModelSlice
 from labelbox.schema.task import Task
 from labelbox.schema.user import User
+from labelbox.schema.editor_task_type import EditorTaskType
 
 logger = logging.getLogger(__name__)
 
@@ -735,6 +736,12 @@ def create_project(self, **kwargs) -> Project:
             } if media_type else {})
         })
 
+    def create_model_chat_project(self, **kwargs) -> Project:
+        kwargs["media_type"] = media_type.MediaType.Conversational
+        kwargs["editor_task_type"] = editor_task_type.EditorTaskType.ModelChatEvaluation
+
+        return self.create_project(**kwargs)
+
     def get_roles(self) -> List[Role]:
         """
         Returns:
@@ -938,10 +945,12 @@ def rootSchemaPayloadToFeatureSchema(client, payload):
                                    rootSchemaPayloadToFeatureSchema,
                                    ['rootSchemaNodes', 'nextCursor'])
 
-    def create_ontology_from_feature_schemas(self,
-                                             name,
-                                             feature_schema_ids,
-                                             media_type=None) -> Ontology:
+    def create_ontology_from_feature_schemas(
+            self,
+            name,
+            feature_schema_ids,
+            media_type: MediaType = None,
+            editor_task_type: EditorTaskType = None) -> Ontology:
         """
         Creates an ontology from a list of feature schema ids
 
@@ -979,7 +988,10 @@ def create_ontology_from_feature_schemas(self,
                     "Neither `tool` or `classification` found in the normalized feature schema"
                 )
         normalized = {'tools': tools, 'classifications': classifications}
-        return self.create_ontology(name, normalized, media_type)
+        return self.create_ontology(name=name,
+                                    normalized=normalized,
+                                    media_type=media_type,
+                                    editor_task_type=editor_task_type)
 
     def delete_unused_feature_schema(self, feature_schema_id: str) -> None:
         """
@@ -1166,7 +1178,11 @@ def get_unused_feature_schemas(self, after: str = None) -> List[str]:
                 "Failed to get unused feature schemas, message: " +
                 str(response.json()['message']))
 
-    def create_ontology(self, name, normalized, media_type=None) -> Ontology:
+    def create_ontology(self,
+                        name,
+                        normalized,
+                        media_type: MediaType = None,
+                        editor_task_type: EditorTaskType = None) -> Ontology:
         """
         Creates an ontology from normalized data
             >>> normalized = {"tools" : [{'tool': 'polygon',  'name': 'cat', 'color': 'black'}], "classifications" : []}
@@ -1194,6 +1210,13 @@ def create_ontology(self, name, normalized, media_type=None) -> Ontology:
             else:
                 raise get_media_type_validation_error(media_type)
 
+        if editor_task_type:
+            if EditorTaskType.is_supported(editor_task_type):
+                editor_task_type = editor_task_type.value
+            else:
+                raise EditorTaskType.get_editor_task_type_validation_error(
+                    editor_task_type)
+
         query_str = """mutation upsertRootSchemaNodePyApi($data:  UpsertOntologyInput!){
                            upsertOntology(data: $data){ %s }
         } """ % query.results_query_part(Entity.Ontology)
@@ -1204,9 +1227,65 @@ def create_ontology(self, name, normalized, media_type=None) -> Ontology:
                 'mediaType': media_type
             }
         }
+        if editor_task_type:
+            params['data']['editorTaskType'] = editor_task_type
+
         res = self.execute(query_str, params)
         return Entity.Ontology(self, res['upsertOntology'])
 
+    def create_model_chat_evaluation_ontology(self, name, normalized):
+        """
+        Creates a model chat evalutation ontology from normalized data
+            >>> normalized = {"tools" : [{'tool': 'message-single-selection', 'name': 'model output single selection', 'color': '#ff0000',},
+                                         {'tool': 'message-multi-selection', 'name': 'model output multi selection', 'color': '#00ff00',},
+                                         {'tool': 'message-ranking', 'name': 'model output multi ranking', 'color': '#0000ff',}]
+                             }
+            >>> ontology = client.create_ontology("ontology-name", normalized)
+
+        Or use the ontology builder
+            >>> ontology_builder = OntologyBuilder(tools=[
+                Tool(tool=Tool.Type.MESSAGE_SINGLE_SELECTION,
+                    name="model output single selection"),
+                Tool(tool=Tool.Type.MESSAGE_MULTI_SELECTION,
+                    name="model output multi selection"),
+                Tool(tool=Tool.Type.MESSAGE_RANKING,
+                    name="model output multi ranking"),
+            ],)
+
+            >>> ontology = client.create_model_chat_evaluation_ontology("Multi-chat ontology", ontology_builder.asdict())
+
+        Args:
+            name (str): Name of the ontology
+            normalized (dict): A normalized ontology payload. See above for details.
+        Returns:
+            The created Ontology
+        """
+
+        return self.create_ontology(
+            name=name,
+            normalized=normalized,
+            media_type=MediaType.Conversational,
+            editor_task_type=EditorTaskType.ModelChatEvaluation)
+
+    def create_model_chat_evaluation_ontology_from_feature_schemas(
+            self, name, feature_schema_ids):
+        """
+        Creates an ontology from a list of feature schema ids
+
+        Args:
+            name (str): Name of the ontology
+            feature_schema_ids (List[str]): List of feature schema ids corresponding to
+                top level tools and classifications to include in the ontology
+        Returns:
+            The created Ontology
+        """
+
+        return self.create_ontology_from_feature_schemas(
+            name=name,
+            feature_schema_ids=feature_schema_ids,
+            media_type=MediaType.Conversational,
+            editor_task_type=EditorTaskType.ModelChatEvaluation)
+
     def create_feature_schema(self, normalized):
         """
         Creates a feature schema from normalized data.

libs/labelbox/src/labelbox/schema/editor_task_type.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+
+class EditorTaskType(Enum):
+    ModelEvaluationWithUploadedAsset = "MODEL_EVALUATION_WITH_UPLOADED_ASSET",
+    ResponseCreation = "RESPONSE_CREATION",
+    ModelChatEvaluation = "MODEL_CHAT_EVALUATION"
+
+    @classmethod
+    def is_supported(cls, value):
+        return isinstance(value, cls)
+
+    @classmethod
+    def get_media_type_validation_error(cls, editor_task_type):
+        return TypeError(f"{editor_task_type}: is not a valid media type. Use"
+                         f" any of {EditorTaskType.__members__.items()}"
+                         " from EditorTaskType.")

libs/labelbox/src/labelbox/schema/ontology.py

@@ -252,6 +252,9 @@ class Type(Enum):
         LINE = "line"
         NER = "named-entity"
         RELATIONSHIP = "edge"
+        MESSAGE_SINGLE_SELECTION = 'message-single-selection'
+        MESSAGE_MULTI_SELECTION = 'message-multi-selection'
+        MESSAGE_RANKING = 'message-ranking'
 
     tool: Type
     name: str

libs/labelbox/tests/integration/conftest.py

@@ -339,6 +339,25 @@ def _upload_invalid_data_rows_for_dataset(dataset: Dataset):
     return _upload_invalid_data_rows_for_dataset
 
 
+@pytest.fixture
+def model_chat_evaluation_ontology(client, rand_gen):
+    ontology_name = f"test-model-chat-evaluation-ontology-{rand_gen(str)}"
+    ontology_builder = OntologyBuilder(tools=[
+        Tool(tool=Tool.Type.MESSAGE_SINGLE_SELECTION,
+             name="model output single selection"),
+        Tool(tool=Tool.Type.MESSAGE_MULTI_SELECTION,
+             name="model output multi selection"),
+        Tool(tool=Tool.Type.MESSAGE_RANKING, name="model output multi ranking"),
+    ],)
+
+    ontology = client.create_model_chat_evaluation_ontology(
+        ontology_name, ontology_builder.asdict())
+
+    yield ontology
+
+    client.delete_unused_ontology(ontology.uid)
+
+
 def pytest_configure():
     pytest.report = defaultdict(int)
 

libs/labelbox/tests/integration/test_model_chat_evaluation_ontology.py

@@ -0,0 +1,87 @@
+import pytest
+from labelbox import OntologyBuilder, Tool
+from labelbox import MediaType
+
+
+def test_create_model_chat_evaluation_ontology(model_chat_evaluation_ontology):
+    ontology = model_chat_evaluation_ontology
+
+    # here we are essentially testing the ontology creation which is a fixture
+    assert ontology
+    assert ontology.name
+    assert len(ontology.tools()) == 3
+    for tool in ontology.tools():
+        assert tool.schema_id
+        assert tool.feature_schema_id
+
+
+@pytest.fixture
+def tools_json():
+    tools = [{
+        'tool': 'message-single-selection',
+        'name': 'model output single selection',
+        'required': False,
+        'color': '#ff0000',
+        'classifications': [],
+        'schemaNodeId': None,
+        'featureSchemaId': None
+    }, {
+        'tool': 'message-multi-selection',
+        'name': 'model output multi selection',
+        'required': False,
+        'color': '#00ff00',
+        'classifications': [],
+        'schemaNodeId': None,
+        'featureSchemaId': None
+    }, {
+        'tool': 'message-ranking',
+        'name': 'model output multi ranking',
+        'required': False,
+        'color': '#0000ff',
+        'classifications': [],
+        'schemaNodeId': None,
+        'featureSchemaId': None
+    }]
+
+    return tools
+
+
+@pytest.fixture
+def features_from_json(client, tools_json):
+    tools = tools_json
+    features = {client.create_feature_schema(t) for t in tools if t}
+
+    yield features
+
+    for f in features:
+        client.delete_unused_feature_schema(f.uid)
+
+
+@pytest.fixture
+def ontology_from_feature_ids(client, features_from_json):
+    feature_ids = {f.uid for f in features_from_json}
+    ontology = client.create_model_chat_evaluation_ontology_from_feature_schemas(
+        name="test-model-chat-evaluation-ontology{rand_gen(str)}",
+        feature_schema_ids=feature_ids,
+    )
+
+    yield ontology
+
+    client.delete_unused_ontology(ontology.uid)
+
+
+def test_ontology_create_feature_schema(ontology_from_feature_ids, features_from_json, tools_json):
+    created_ontology = ontology_from_feature_ids
+    feature_schema_ids = {f.uid for f in features_from_json}
+    tools_normalized = created_ontology.normalized['tools']
+    tools = tools_json
+
+    for tool in tools:
+        generated_tool = next(
+            t for t in tools_normalized if t['name'] == tool['name'])
+        assert generated_tool['schemaNodeId'] is not None
+        assert generated_tool['featureSchemaId'] in feature_schema_ids
+        assert generated_tool['tool'] == tool['tool']
+        assert generated_tool['name'] == tool['name']
+        assert generated_tool['required'] == tool['required']
+        assert generated_tool['color'] == tool['color']