Update CreateExport query (#162)

nateshim-indico · Nathanael Shim · web-flow · commit a09d4c157f6e · 2022-05-10T07:30:36.000-07:00
* updating createExport query

* addresses comments; runs sphinx-build to make sure documentation is updated

Co-authored-by: Nathanael Shim &lt;nathanaelshim@Nathanaels-MacBook-Pro.local&gt;
diff --git a/docs/exports.html b/docs/exports.html
@@ -141,17 +141,18 @@
 
 <dl class="py class">
 <dt class="sig sig-object py" id="indico.queries.export.CreateExport">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">indico.queries.export.</span></span><span class="sig-name descname"><span class="pre">CreateExport</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">dataset_id</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">subset_ids</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">column_ids</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">labelset_ids</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">file_info</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">combine_labels</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">anonymous</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">wait</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#indico.queries.export.CreateExport" title="Permalink to this definition"></a></dt>
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">indico.queries.export.</span></span><span class="sig-name descname"><span class="pre">CreateExport</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">dataset_id</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">labelset_id</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">column_ids</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">model_ids</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">frozen_labelset_ids</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">combine_labels</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">'ALL'</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">file_info</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">anonymous</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">wait</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#indico.queries.export.CreateExport" title="Permalink to this definition"></a></dt>
 <dd><p>Create an export job for a dataset.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>dataset_id</strong> (<em>int</em>) – Dataset to create the export for</p></li>
-<li><p><strong>subset_ids</strong> (<code class="xref py py-data docutils literal notranslate"><span class="pre">Optional</span></code>[<code class="xref py py-class docutils literal notranslate"><span class="pre">List</span></code>[<code class="xref py py-class docutils literal notranslate"><span class="pre">int</span></code>]]) – (List(int)): Subset ids to export rows sets</p></li>
+<li><p><strong>labelset_id</strong> (<code class="xref py py-class docutils literal notranslate"><span class="pre">int</span></code>) – (int): Labelset column id to export</p></li>
 <li><p><strong>column_ids</strong> (<em>List</em><em>(</em><em>int</em><em>)</em>) – Data column ids to export</p></li>
-<li><p><strong>labelset_ids</strong> (<code class="xref py py-data docutils literal notranslate"><span class="pre">Optional</span></code>[<code class="xref py py-class docutils literal notranslate"><span class="pre">List</span></code>[<code class="xref py py-class docutils literal notranslate"><span class="pre">int</span></code>]]) – (List(int)): Labelset column ids to export</p></li>
+<li><p><strong>model_ids</strong> (<em>List</em><em>(</em><em>int</em><em>)</em>) – Model ids to include predictions from</p></li>
+<li><p><strong>frozen_labelset_ids</strong> (<code class="xref py py-data docutils literal notranslate"><span class="pre">Optional</span></code>[<code class="xref py py-class docutils literal notranslate"><span class="pre">List</span></code>[<code class="xref py py-class docutils literal notranslate"><span class="pre">int</span></code>]]) – (List(int)): frozen labelset ids to limit examples by</p></li>
+<li><p><strong>combine_labels</strong> (<code class="xref py py-class docutils literal notranslate"><span class="pre">LabelResolutionStrategy</span></code>) – (LabelResolutionStrategy): One row per example, combine labels from multiple labels into a single row</p></li>
 <li><p><strong>file_info</strong> (<code class="xref py py-class docutils literal notranslate"><span class="pre">bool</span></code>) – (bool): Include datafile information</p></li>
-<li><p><strong>combine_labels</strong> (<code class="xref py py-class docutils literal notranslate"><span class="pre">bool</span></code>) – (bool): One row per example, combine labels from multiple labels into a single row</p></li>
 <li><p><strong>anonymous</strong> (<code class="xref py py-class docutils literal notranslate"><span class="pre">bool</span></code>) – (bool): Anonymize user information</p></li>
 <li><p><strong>wait</strong> (<code class="xref py py-class docutils literal notranslate"><span class="pre">bool</span></code>) – (bool): Wait for the export to complete. Default is True</p></li>
 </ul>
diff --git a/docs/searchindex.js b/docs/searchindex.js
diff --git a/examples/create-download-export.py b/examples/create-download-export.py
@@ -0,0 +1,29 @@
+from indico import IndicoClient, IndicoConfig
+from indico.queries import GetDataset, CreateExport, DownloadExport
+import pandas as pd
+
+
+"""Example 1: Create export from dataset and download as csv"""
+
+dataset_id = 6826
+
+# Create an Indico API client
+my_config = IndicoConfig(
+    host="app.indico.io", api_token_path="./path/to/indico_api_token.txt"
+)
+client = IndicoClient(config=my_config)
+
+# Get dataset object
+dataset = client.call(GetDataset(id=dataset_id))
+
+# Create export object using dataset's id and labelset id
+export = client.call(
+            CreateExport(
+                dataset_id=dataset.id,
+                labelset_id=dataset.labelsets[0].id,
+                wait=True
+            )
+        )
+
+# Use export object to download as pandas csv
+csv = client.call(DownloadExport(export.id))
diff --git a/indico/queries/export.py b/indico/queries/export.py
@@ -4,63 +4,68 @@
 
 from indico.client import GraphQLRequest, RequestChain, Debouncer
 from indico.errors import IndicoNotFound, IndicoRequestError
-from indico.types.export import Export
+from indico.types.export import LabelResolutionStrategy, Export
 from indico.queries.storage import RetrieveStorageObject
 
 
 class _CreateExport(GraphQLRequest):
     query = """
         mutation CreateExport(
-            $datasetId: Int!, 
+            $datasetId: Int!,
+            $labelsetId: Int!,
             $columnIds: [Int], 
-            $subsetIds: [Int], 
-            $labelsetIds: [Int],
+            $modelIds: [Int],
+            $frozenLabelsetIds: [Int],
+            $combineLabels: LabelResolutionStrategy,
             $fileInfo: Boolean,
-            $combineLabels: Boolean,
             $anonymous: Boolean
         ) {
             createExport(
                 datasetId: $datasetId,
+                labelsetId: $labelsetId,
                 columnIds: $columnIds,
-                subsetIds: $subsetIds,
+                modelIds: $modelIds,
+                frozenLabelsetIds: $frozenLabelsetIds,
+                combineLabels: $combineLabels,
                 fileInfo: $fileInfo,
-                labelsetIds: $labelsetIds,
-                combineLabels: $combineLabels
                 anonymous: $anonymous
             ) {
                 id,
                 datasetId,
+                labelsetId,
                 name,
                 status,
-                downloadUrl,
                 columnIds,
-                labelsetIds,
-                subsetIds,
+                modelIds,
+                frozenLabelsetIds,
                 anonymous
+                downloadUrl,
             }
-        }    
+        }
 
     """
 
     def __init__(
         self,
         dataset_id: int,
+        labelset_id: int,
         column_ids: List[int] = None,
-        subset_ids: List[int] = None,
-        labelset_ids: List[int] = None,
+        model_ids: List[int] = None,
+        frozen_labelset_ids: List[int] = None,
+        combine_labels: LabelResolutionStrategy = LabelResolutionStrategy.ALL.name,
         file_info: bool = None,
-        combine_labels: bool = None,
         anonymoous: bool = None,
     ):
         super().__init__(
             self.query,
             variables={
                 "datasetId": dataset_id,
+                "labelsetId": labelset_id,
                 "columnIds": column_ids,
-                "subsetIds": subset_ids,
-                "labelsetIds": labelset_ids,
-                "fileInfo": file_info,
+                "modelIds": model_ids,
+                "frozenLabelsetIds": frozen_labelset_ids,
                 "combineLabels": combine_labels,
+                "fileInfo": file_info,
                 "anonymous": anonymoous,
             },
         )
@@ -91,8 +96,9 @@ class GetExport(GraphQLRequest):
                 name 
                 status
                 columnIds
-                labelsetIds
-                subsetIds
+                labelsetId
+                modelIds
+                frozenLabelsetIds
                 numLabels
                 anonymous
                 downloadUrl
@@ -160,11 +166,12 @@ class CreateExport(RequestChain):
 
     Args:
         dataset_id (int): Dataset to create the export for
-        subset_ids: (List(int)): Subset ids to export rows sets
+        labelset_id: (int): Labelset column id to export
         column_ids (List(int)): Data column ids to export
-        labelset_ids: (List(int)): Labelset column ids to export
+        model_ids (List(int)): Model ids to include predictions from
+        frozen_labelset_ids: (List(int)): frozen labelset ids to limit examples by
+        combine_labels: (LabelResolutionStrategy): One row per example, combine labels from multiple labels into a single row
         file_info: (bool): Include datafile information
-        combine_labels: (bool): One row per example, combine labels from multiple labels into a single row
         anonymous: (bool): Anonymize user information
         wait: (bool): Wait for the export to complete. Default is True
 
@@ -178,32 +185,35 @@ class CreateExport(RequestChain):
     def __init__(
         self,
         dataset_id: int,
-        subset_ids: List[int] = None,
+        labelset_id: int,
         column_ids: List[int] = None,
-        labelset_ids: List[int] = None,
+        model_ids: List[int] = None,
+        frozen_labelset_ids: List[int] = None,
+        combine_labels: LabelResolutionStrategy = LabelResolutionStrategy.ALL.name,
         file_info: bool = False,
-        combine_labels: bool = False,
         anonymous: bool = False,
         wait: bool = True,
     ):
         self.dataset_id = dataset_id
+        self.labelset_id = labelset_id
         self.column_ids = column_ids
-        self.subset_ids = subset_ids
-        self.labelset_ids = labelset_ids
-        self.file_info = file_info
+        self.model_ids = model_ids
+        self.frozen_labelset_ids = frozen_labelset_ids
         self.combine_labels = combine_labels
+        self.file_info = file_info
         self.anonymous = anonymous
         self.wait = wait
         super().__init__()
 
     def requests(self):
         yield _CreateExport(
             dataset_id=self.dataset_id,
+            labelset_id=self.labelset_id,
             column_ids=self.column_ids,
-            subset_ids=self.subset_ids,
-            labelset_ids=self.labelset_ids,
-            file_info=self.file_info,
+            model_ids=self.model_ids,
+            frozen_labelset_ids=self.frozen_labelset_ids,
             combine_labels=self.combine_labels,
+            file_info=self.file_info,
             anonymoous=self.anonymous,
         )
         debouncer = Debouncer()
diff --git a/indico/types/export.py b/indico/types/export.py
@@ -1,19 +1,29 @@
 from typing import List
+from enum import Enum
 
 from indico.types.base import BaseType
 
 
+class LabelResolutionStrategy(Enum):
+    MAJORITY_VOTE_WITH_TIES = "majority_vote_with_ties"
+    MAJORITY_VOTE_WITHOUT_TIES = "majority_vote_without_ties"
+    UNANIMOUS = "unanimous"
+    ALL = "all"
+
+
 class Export(BaseType):
     """
     An export of a dataset.
 
     Attributes:
         id (int): Id of the export
         dataset_id (int): Dataset id of the export
-        status (str): Export job status
+        name (str): Name of the export
+        status (ExportStatus): Export job status
         column_ids (List(int)): Data columns for export
-        labelset_ids (List(int)): Labelsets columns for export
-        subset_ids (List(int)): Subsets requested for export rows
+        labelset_id (int): Labelset column for export
+        model_ids (List(int)): Models to include predictions from
+        frozen_labelset_ids (List(int)): Frozen labelsets to select examples
         num_labels (int): Number of labels on the dataset
         anonymous (bool): Whether to anonymize labelers
         download_url (str): Indico url of the export csv
@@ -25,8 +35,9 @@ class Export(BaseType):
     name: str
     status: str
     column_ids: List[int]
-    labelset_ids: List[int]
-    subset_ids: List[int]
+    labelset_id: int
+    model_ids: List[int]
+    frozen_labelset_ids: List[int]
     num_labels: int
     anonymous: bool
     download_url: str
diff --git a/tests/integration/queries/test_dataset.py b/tests/integration/queries/test_dataset.py
@@ -287,7 +287,12 @@ def test_create_from_csv(indico):
 
     _dataset_complete(dataset)
 
-    export = client.call(CreateExport(dataset_id=dataset.id, wait=True))
+    export = client.call(CreateExport(
+                            dataset_id=dataset.id,
+                            labelset_id=dataset.labelsets[0].id,
+                            wait=True
+                        )
+                    )
 
     exported_data = client.call(DownloadExport(export.id))
 
diff --git a/tests/integration/queries/test_export.py b/tests/integration/queries/test_export.py
@@ -11,11 +11,15 @@
 
 def test_create_and_download_export(airlines_dataset: Dataset):
     client = IndicoClient()
-    export = client.call(CreateExport(dataset_id=airlines_dataset.id, wait=True))
+    export = client.call(
+                CreateExport(
+                    dataset_id=airlines_dataset.id,
+                    labelset_id=airlines_dataset.labelsets[0].id,
+                    wait=True
+                )
+            )
     assert export.status == "COMPLETE"
-
     csv = client.call(DownloadExport(export.id))
-
     assert all(c in csv.columns.to_list() for c in ["ID", "Target_1", "Text"])
     assert any(re.match("row_index_[0-9]+", c) for c in csv.columns.to_list())
     assert csv["Text"][0] == "Your service is so bad."
@@ -35,5 +39,10 @@ def test_download_incomplete(indico):
 
 def test_create_export_no_wait(airlines_dataset: Dataset):
     client = IndicoClient()
-    export = client.call(CreateExport(dataset_id=airlines_dataset.id, wait=False))
+    export = client.call(CreateExport(
+                            dataset_id=airlines_dataset.id,
+                            labelset_id=airlines_dataset.labelsets[0].id,
+                            wait=False
+                        )
+                    )
     assert export.status == "STARTED"
