indico_toolkit 6.14.0__py3-none-any.whl

indico_toolkit/indico_wrapper/download.py

@@ -0,0 +1,193 @@
+import os
+from io import StringIO
+
+from indico import IndicoClient, IndicoRequestError
+from indico.queries import (
+    CreateExport,
+    DownloadExport,
+    GraphQLRequest,
+    RetrieveStorageObject,
+)
+from indico.types.export import Export
+
+from ..errors import ToolkitInputError
+from ..retry import retry
+
+try:
+    import pandas as pd
+
+    _PANDAS_INSTALLED = True
+except ImportError as error:
+    _PANDAS_INSTALLED = False
+    _IMPORT_ERROR = error
+
+
+class Download:
+    """
+    Class to support downloading resources from an Indico Cluster
+    """
+
+    def __init__(self, client: IndicoClient):
+        self.client = client
+
+    def get_dataset_pdfs(
+        self,
+        dataset_id: int,
+        labelset_id: int,
+        output_dir: str,
+        max_files_to_download: int = None,
+    ) -> int:
+        """Download PDFs from an uploaded dataset to a local directory
+
+        Args:
+            dataset_id (int): Dataset ID to download from
+            labelset_id (int): ID of your labelset (from teach task)
+            output_dir (str): Path to directory to write PDFs
+            max_files_to_download (int): = Max number of files to download
+                (default: None = download all)
+
+        Raises:
+            ToolkitInputError: Exception if invalid directory path
+
+        Returns:
+            int: Number of files downloaded
+        """
+        if not os.path.isdir(output_dir):
+            raise ToolkitInputError(f"Path is not a directory: {output_dir}")
+        export = self._create_export(dataset_id, labelset_id)
+        df = self._download_export(export.id)
+        num_files_downloaded = self._download_pdfs_from_export(
+            df,
+            output_dir,
+            f"file_name_{dataset_id}",
+            f"file_url_{dataset_id}",
+            max_files_to_download,
+        )
+        return num_files_downloaded
+
+    def get_uploaded_csv_dataframe(self, dataset_id: int) -> "pd.DataFrame":
+        """
+        Get a dataframe from a CSV that has been uploaded to the platform
+        Args:
+            dataset_id (int): id of the dataset
+        Returns:
+            pd.DataFrame: a dataframe representation of the CSV you uploaded
+        """
+        if not _PANDAS_INSTALLED:
+            raise RuntimeError(
+                "getting an uploaded CSV dataframe requires additional dependencies: "
+                "`pip install indico-toolkit[downloads]`"
+            ) from _IMPORT_ERROR
+
+        url = self._get_csv_download_url(dataset_id)
+        string_df = self._retrieve_storage_object(url)
+        return pd.read_csv(StringIO(string_df))
+
+    def get_snapshot_dataframe(
+        self, dataset_id: int, labelset_id: int, file_info: bool = True, **kwargs
+    ) -> "pd.DataFrame":
+        """Download a snapshot. For additional arguments,
+        see documentation for CreateExport in the Python SDK.
+
+        Args:
+            dataset_id (int): dataset ID you're interested in
+            labelset_id (int): ID of your labelset (from teach task)
+            file_info (bool, optional): Include additional file level metadata.
+                Defaults to True.
+
+        Returns:
+            pd.DataFrame: DataFrame with full document text and additional metadata
+        """
+        export = self._create_export(
+            dataset_id, labelset_id, file_info=file_info, **kwargs
+        )
+        return self._download_export(export.id)
+
+    def _download_pdfs_from_export(
+        self,
+        export_df: "pd.DataFrame",
+        output_dir: str,
+        file_name_col: str,
+        file_url_col: str,
+        max_files_to_download: int = None,
+    ) -> int:
+        try:
+            import tqdm
+        except ImportError as error:
+            raise RuntimeError(
+                "downloading pdfs requires additional dependencies: "
+                "`pip install indico-toolkit[downloads]`"
+            ) from error
+
+        for i, row in tqdm.tqdm(export_df.iterrows()):
+            basename = os.path.basename(row[file_name_col])
+            pdf_bytes = self._retrieve_storage_object(row[file_url_col])
+            with open(os.path.join(output_dir, basename), "wb") as fd:
+                fd.write(pdf_bytes)
+            if max_files_to_download and i + 1 == max_files_to_download:
+                return max_files_to_download
+        return export_df.shape[0]
+
+    @retry(IndicoRequestError, ConnectionError)
+    def _download_export(self, export_id: int) -> "pd.DataFrame":
+        """
+        Download a dataframe representation of your dataset export
+        """
+        return self.client.call(DownloadExport(export_id=export_id))
+
+    @retry(IndicoRequestError, ConnectionError)
+    def _create_export(
+        self,
+        dataset_id: int,
+        labelset_id: int,
+        file_info: bool = True,
+        wait: bool = True,
+        **kwargs,
+    ) -> Export:
+        """
+        Create an export object that can be used to download datasets. For additional
+        kwargs, see CreateExport docstring in the Python SDK.
+
+        Args:
+            dataset_id (int): ID of your dataset
+            labelset_id (int): ID of your labelset (from teach task)
+            file_info (bool, optional): whether to include additional file metadata.
+                Defaults to True.
+            wait (bool, optional): wait for export to be created. Defaults to True.
+
+        Returns:
+            Export: Description of dataset assets. See Python SDK for full object
+                description
+        """
+        return self.client.call(
+            CreateExport(
+                dataset_id=dataset_id,
+                file_info=file_info,
+                labelset_id=labelset_id,
+                wait=wait,
+                **kwargs,
+            )
+        )
+
+    @retry(IndicoRequestError, ConnectionError)
+    def _retrieve_storage_object(self, url: str):
+        return self.client.call(RetrieveStorageObject(url))
+
+    def _get_csv_download_url(self, dataset_id: int) -> str:
+        query = """
+            query getDataURL($id: Int!) {
+                dataset(id: $id) {
+                    files {
+                    fileType
+                    id
+                    rainbowUrl
+                    }
+                }
+            }
+        """
+        result = self.client.call(GraphQLRequest(query, {"id": dataset_id}))
+        # loop through in case there are other file types uploaded to dataset
+        for file in result["dataset"]["files"]:
+            if file["fileType"] == "CSV":
+                return file["rainbowUrl"]
+        raise ToolkitInputError(f"There are no CSVs uploaded to {dataset_id}")

indico_toolkit/indico_wrapper/indico_wrapper.py

@@ -0,0 +1,119 @@
+from typing import List, Union
+
+from indico import IndicoClient
+from indico.errors import IndicoRequestError
+from indico.queries import (
+    AddModelGroupComponent,
+    CreateStorageURLs,
+    GraphQLRequest,
+    JobStatus,
+    ModelGroupPredict,
+    RetrieveStorageObject,
+)
+from indico.types import Dataset, ModelGroup, Workflow
+
+from ..errors import ToolkitStatusError
+from ..retry import retry
+from ..types import Predictions
+
+
+class IndicoWrapper:
+    """
+    Class for shared API functionality
+    """
+
+    def __init__(self, client: IndicoClient):
+        """
+        Create indico client with user provided arguments
+
+        args:
+            client (IndicoClient): instantiated Indico Client object
+        """
+        self.client = client
+
+    def train_model(
+        self,
+        dataset: Dataset,
+        workflow: Workflow,
+        model_name: str,
+        source_col: str,
+        target_col: str,
+        after_component_id: int = None,
+    ) -> ModelGroup:
+        """
+        Train an Indico model
+        Args:
+            dataset (Dataset): A dataset object (should represent an uploaded CSV
+                dataset)
+            workflow (Workflow): A workflow object to the corresponding dataset
+            model_name (str): the name for your model
+            source_col (str): the csv column that contained the text
+            target_col (str): the csv column that contained the labels
+            after_component_id (int, optional): The workflow component that precedes
+                this model group. If None, will be set programmatically to the id of the
+                Input OCR Extraction component. Defaults to None.
+            wait (bool, optional): Wait for the model to finish training.
+                Defaults to False.
+
+        Returns:
+            ModelGroup: Model group object
+        """
+        if not after_component_id:
+            after_component_id = workflow.component_by_type("INPUT_OCR_EXTRACTION").id
+        return self.client.call(
+            AddModelGroupComponent(
+                name=model_name,
+                dataset_id=dataset.id,
+                source_column_id=dataset.datacolumn_by_name(source_col).id,
+                labelset_column_id=dataset.labelset_by_name(target_col).id,
+                workflow_id=workflow.id,
+                after_component_id=after_component_id,
+            )
+        )
+
+    @retry(IndicoRequestError, ConnectionError)
+    def get_storage_object(self, storage_url: str):
+        return self.client.call(RetrieveStorageObject(storage_url))
+
+    def create_storage_urls(self, file_paths: List[str]) -> List[str]:
+        return self.client.call(CreateStorageURLs(files=file_paths))
+
+    def get_job_status(self, job_id: int, wait: bool = True, timeout: float = None):
+        return self.client.call(JobStatus(id=job_id, wait=wait, timeout=timeout))
+
+    @retry(IndicoRequestError, ConnectionError)
+    def graphQL_request(self, graphql_query: str, variables: dict = None):
+        return self.client.call(
+            GraphQLRequest(query=graphql_query, variables=variables)
+        )
+
+    def get_predictions_with_model_id(
+        self,
+        model_id: int,
+        samples: List[str],
+        load: bool = False,
+        options: dict = None,
+        wait: bool = True,
+    ) -> Union[int, List[Predictions]]:
+        """
+        Submit samples directly to a model. Note: documents must already be in raw text.
+        Args:
+            model_id (int): The model ID to submit to
+            samples (List[str]): A list containing the text samples you want to submit
+            load (bool, optional): [deprecated] Always set to False. Defaults to False.
+            options (dict, optional): Model Prediction options. Defaults to None.
+            wait (bool, optional): Wait for predictions to finish. Defaults to True.
+
+        Returns: if wait is False, returns the job ID, else returns a list of
+            Predictions where each Predictions is either type Classifications or
+            Extractions depending on your model.
+        """
+        job = self.client.call(ModelGroupPredict(model_id, samples, load, options))
+        if not wait:
+            return job.id
+        status = self.get_job_status(job.id, wait=True)
+        if status.status != "SUCCESS":
+            raise ToolkitStatusError(
+                f"Predictions Failed, {status.status}: {status.result}"
+            )
+        return [Predictions.get_obj(i) for i in status.result]

indico_toolkit/indico_wrapper/modelop.py

@@ -0,0 +1,200 @@
+import json
+from typing import Iterator
+
+from indico import IndicoClient
+from indico.queries import GraphQLRequest
+
+# valid model option parameters
+TEXT_EXTRACTION_PARAMS = {
+    "max_empty_chunk_ratio": lambda value: 0 <= value <= 1.0e5,
+    "auto_negative_scaling": [True, False],
+    "optimize_for": [
+        "predict_speed",
+        "accuracy",
+        "speed",
+        "accuracy_fp16",
+        "predict_speed_fp16",
+    ],
+    "subtoken_predictions": [True, False],
+    "base_model": [
+        "roberta",
+        "small",
+        "multilingual",
+        "fast",
+        "textcnn",
+        "fasttextcnn",
+    ],
+    "class_weight": ["linear", "sqrt", "log", None],
+}
+
+TEXT_CLASSIFICATION_PARAMS = {
+    "model_type": ["tfidf_lr", "tfidf_gbt", "standard", "finetune"]
+}
+
+
+class ModelOp:
+    """
+    Class to assit in Model Group Operation calls.
+    """
+
+    def __init__(self, client: IndicoClient):
+        self.client = client
+
+    def get_model_options(
+        self, model_group_id: int, model_id: int | None = None
+    ) -> dict[str, object]:
+        """
+        Return model options for `model_id` of `model_group_id`
+        or the most recent model if `model_id` is not specified.
+        Args:
+            model_group_id (int): id of model group
+            model_id (int, optional): argument to return a specific model within a model
+                group
+        """
+        all_model_options = self.get_all_model_options(model_group_id)
+
+        if model_id is None:
+            return next(all_model_options)
+
+        for model_options in all_model_options:
+            if model_options["id"] == model_id:
+                return model_options
+        else:
+            raise RuntimeError(
+                f"Model group {model_group_id} does not have a model with ID {model_id}"
+            )
+
+    def get_all_model_options(self, model_group_id: int) -> Iterator[dict[str, object]]:
+        """
+        Return model options for all models of `model_group_id`.
+        Args:
+            model_group_id (int): id of model group
+        Returns:
+            dict: dictionary of all model options, including:
+            "id", "domain, ""high_quality", "interlabeler_resolution",
+            "sampling_strategy", "seed", "test_split", "weight_by_class_frequency",
+            "word_predictor_strength", and "model_training_options"
+        """
+
+        response = self.client.call(
+            GraphQLRequest(
+                """
+                query getMoonbowModelGroupModels($modelGroupId: Int!) {
+                    modelGroup(modelGroupId: $modelGroupId) {
+                        id
+                        models {
+                            id
+                            modelOptions
+                        }
+                    }
+                }
+                """,
+                {"modelGroupId": model_group_id},
+            )
+        )
+
+        for model in response["modelGroup"]["models"]:
+            model_options = json.loads(model["modelOptions"])
+            model_options["id"] = model["id"]
+            yield model_options
+
+    def retrain_model(self, model_group_id: int) -> dict:
+        """
+        Forces model retrain.
+        Args:
+            model_group_id (int): id of model group needing retraining
+        Returns:
+            dict: returns model group id and status of retraining model
+        """
+        return self.client.call(
+            GraphQLRequest(
+                """
+                mutation retrainMoonbowModelGroup($modelGroupId: Int!) {
+                    retrainModelGroup(modelGroupId: $modelGroupId, forceRetrain: true) {
+                        id
+                        status
+                    }
+                }    
+                """,
+                {"modelGroupId": model_group_id},
+            )
+        )
+
+    def update_model_settings(
+        self, model_group_id: int, model_type: str, **model_training_options
+    ) -> dict[str, object]:
+        """
+        Update group model settings based on specified model training options.
+
+        Args:
+             model_group_id (int): id of model group
+             model_type (str): type of model ("text_extraction", "text_classification")
+             **kwargs: Advanced Model Training Options
+
+            Default values of Advanced Model Training Options and avaiable parameters:
+                    For Text Extraction Model:
+                        max_empty_chunk_ratio : 1.0 (min of 0, no max value: a large
+                            number effectivly turns this option off)
+                        auto_negative_scaling : True
+                        optimize_for : "predict_speed" ( "predict_speed", "accuracy",
+                            "speed", "accuracy_fp16" and "predict_speed_fp16")
+                        subtoken_predictions : True
+                        base_model : "roberta" ("roberta", "small" (distilled version of
+                            RoBERTa), "multilingual", "fast", "textcnn", "fasttextcnn")
+                        class_weight : "sqrt" ("linear", "sqrt", "log", None)
+
+                    For Text Classification Model:
+                        model_type : "standard" ("tfidf_lr", "tfidf_gbt", "standard",
+                            "finetune")
+        Returns:
+            dict: Dictionary of advanced model training options
+        """
+        model = self.client.call(
+            GraphQLRequest(
+                """
+                mutation updateModelGroup(
+                    $modelGroupId: Int!,
+                    $modelTrainingOptions: JSONString
+                ) {
+                    updateModelGroupSettings(
+                      modelGroupId: $modelGroupId
+                      modelTrainingOptions: $modelTrainingOptions
+                  ) {
+                      modelOptions {
+                          id
+                          modelTrainingOptions
+                        }
+                    }
+                }
+                """,
+                {
+                    "modelGroupId": model_group_id,
+                    "modelTrainingOptions": json.dumps(
+                        self._parameter_check(model_type, **model_training_options)
+                    ),
+                },
+            )
+        )
+        options = json.loads(
+            model["updateModelGroupSettings"]["modelOptions"]["modelTrainingOptions"]
+        )
+        options["id"] = model["updateModelGroupSettings"]["modelOptions"]["id"]
+        return options
+
+    def _parameter_check(self, model_type, **model_training_options):
+        if model_type == "text_extraction":
+            valid_params = TEXT_EXTRACTION_PARAMS
+        elif model_type == "text_classification":
+            valid_params = TEXT_CLASSIFICATION_PARAMS
+        else:
+            raise ValueError(f"Invalid model type: {model_type}")
+
+        for param, value in model_training_options.items():
+            if param not in valid_params:
+                raise ValueError(f"Invalid parameter for {model_type}: {param}")
+            if callable(valid_params[param]) and not valid_params[param](value):
+                raise ValueError(f"Invalid value for parameter '{param}': {value}")
+            if value not in valid_params[param]:
+                raise ValueError(f"Invalid value for parameter '{param}': {value}")
+
+        return model_training_options

indico_toolkit/indico_wrapper/reviewer.py

@@ -0,0 +1,111 @@
+import json
+
+from indico import IndicoClient
+
+from .indico_wrapper import Workflow
+
+
+class Reviewer(Workflow):
+    """
+    Class to simulate human reviewer
+    """
+
+    def __init__(
+        self,
+        client: IndicoClient,
+        workflow_id: int,
+    ):
+        self.client = client
+        self.workflow_id = workflow_id
+
+    def accept_review(self, submission_id: int, changes: dict) -> None:
+        """
+        Accept a submission in the review queue
+        Args:
+            submission_id (int): submission ID
+            changes (dict): accepted predictions with format like,
+                e.g. {"model_name": [{"label"...}]}
+        """
+        self.graphQL_request(
+            SUBMIT_REVIEW,
+            {
+                "rejected": False,
+                "submissionId": submission_id,
+                "changes": json.dumps(changes),
+            },
+        )
+
+    def get_random_review_id(self):
+        response = self.graphQL_request(
+            GET_RANDOM_REVIEW_ID, {"workflowId": self.workflow_id}
+        )
+        try:
+            return response["randomSubmission"]["id"]
+        except Exception:
+            raise RuntimeError("The review queue is empty")
+
+    def get_random_exception_id(self):
+        response = self.graphQL_request(
+            GET_RANDOM_EXCEPTION_ID, {"workflowId": self.workflow_id}
+        )
+        try:
+            return response["randomSubmission"]["id"]
+        except Exception:
+            raise RuntimeError("The exception queue is empty")
+
+    def reject_submission(self, submission_id):
+        return self.graphQL_request(
+            SUBMIT_REVIEW, {"rejected": True, "submissionId": submission_id}
+        )
+
+
+SUBMIT_REVIEW = """
+mutation submitStandardQueue(
+    $changes: JSONString,
+    $rejected: Boolean,
+    $submissionId: Int!,
+    $notes: String
+) {
+  submitReview(
+    changes: $changes,
+    rejected: $rejected,
+    submissionId: $submissionId,
+    notes: $notes
+) {
+    id
+    __typename
+  }
+}
+"""
+
+GET_RANDOM_EXCEPTION_ID = """
+query getExceptionsSubmission($workflowId: Int!) {
+  randomSubmission(adminReview: true, workflowId: $workflowId) {
+    id
+    resultFile
+    inputFilename
+    autoReview {
+      id
+      changes
+      __typename
+    }
+    __typename
+  }
+}
+"""
+
+GET_RANDOM_REVIEW_ID = """
+query getSubmission($workflowId: Int!) {
+  randomSubmission(adminReview: false, workflowId: $workflowId) {
+    id
+    resultFile
+    inputFilename
+    autoReview {
+      id
+      changes
+      __typename
+    }
+    __typename
+  }
+}
+"""