indico_toolkit 6.14.0__py3-none-any.whl

indico_toolkit/snapshots/snapshot.py

@@ -0,0 +1,340 @@
+import json
+import os
+from json import JSONDecodeError
+from typing import List, Tuple, Union
+
+from ..errors import ToolkitInputError, ToolkitInstantiationError
+
+try:
+    import pandas as pd
+
+    _PANDAS_INSTALLED = True
+except ImportError as error:
+    _PANDAS_INSTALLED = False
+    _IMPORT_ERROR = error
+
+# TODO: add functionality for classification snapshots
+
+
+class Snapshot:
+    def __init__(
+        self,
+        path_to_snapshot: str,
+        text_col: str = None,
+        label_col: str = None,
+        file_name_col: str = None,
+    ):
+        """
+        Combine and manipulate a Teach Task's snapshot.
+
+        Args:
+            path_to_snapshot (str): path to Snapshot CSV
+            text_col (str, optional): Column with text, will be inferred if not
+                provided. Defaults to None.
+            label_col (str, optional): Column with labels, will be inferred if not
+                provided. Defaults to None.
+            file_name_col (str, optional): Column with file names, will be inferred if
+                not provided. Defaults to None.
+        """
+        if not _PANDAS_INSTALLED:
+            raise RuntimeError(
+                "snapshots require additional dependencies: "
+                "`pip install indico-toolkit[snapshots]`"
+            ) from _IMPORT_ERROR
+
+        self.path_to_snapshot = path_to_snapshot
+        self.df: pd.DataFrame = pd.read_csv(path_to_snapshot)
+        self.label_col = label_col
+        if label_col is None:
+            self._infer_label_col()
+        self._convert_col_from_json()
+        self.text_col = text_col
+        if text_col is None:
+            self._infer_text_col()
+        self.file_name_col = file_name_col
+        if file_name_col is None:
+            self._infer_file_name_col()
+
+    def remove_extraction_labels(self, labels_to_remove: List[str]):
+        """
+        Remove unwanted labels from label column
+        Args:
+            labels_to_remove (List[str]): the labels you'd like to remove
+        """
+        updated_predictions = []
+        for label_set in self.df[self.label_col]:
+            label_set.update(
+                {
+                    "targets": [
+                        i
+                        for i in label_set["targets"]
+                        if i["label"] not in labels_to_remove
+                    ]
+                }
+            )
+            updated_predictions.append(label_set)
+        self.df[self.label_col] = updated_predictions
+
+    def standardize_column_names(
+        self,
+        text_col: str = "source",
+        label_col: str = "target",
+        file_name_col: str = "file_name",
+    ):
+        """
+        To allow merging/appending snapshots, standardize the key column headers
+        """
+        self.df.rename(
+            mapper={
+                self.file_name_col: file_name_col,
+                self.text_col: text_col,
+                self.label_col: label_col,
+            },
+            inplace=True,
+            errors="raise",
+            axis=1,
+        )
+        self.file_name_col = file_name_col
+        self.label_col = label_col
+        self.text_col = text_col
+
+    def drop_unneeded_columns(self, columns_to_drop: List[str] = None):
+        """
+        Keep only text, labels, and file name columns or specify columns to drop by
+        passing them in as a list.
+        """
+        if columns_to_drop:
+            self.df.drop(labels=columns_to_drop, axis=1, inplace=True)
+        else:
+            self.df = self.df[[self.label_col, self.text_col, self.file_name_col]]
+
+    def append(self, snap_to_add: "Snapshot"):
+        """
+        Append the rows from another Snapshot to this snapshot. Ensure column names are
+        standardized beforehand.
+
+        Args:
+            snap_to_add (Snapshot): Snapshot to add
+        """
+        self._assert_key_column_names_match(snap_to_add)
+        self.df = pd.concat([self.df, snap_to_add.df], ignore_index=True)
+
+    def get_extraction_label_names(self):
+        """
+        Return a list of all labeled classes in an extraction snapshot.
+        """
+        label_column = self.df[self.label_col].tolist()
+        label_set = set()
+        for labels in label_column:
+            for label in labels["targets"]:
+                label_set.add(label["label"])
+        return sorted(list(label_set))
+
+    def merge_by_file_name(
+        self,
+        snap_to_merge: "Snapshot",
+        ensure_identical_text: bool = True,
+    ):
+        """
+        Merge extraction labels for identical files. Merge is 'left' and file names /
+        rows only present in 'snap_to_merge' are excluded.
+
+        Args:
+            snap_to_merge (Snapshot): Snapshot you want to merge
+            ensure_identical_text (bool, optional): Require document text to be
+                identical for common file name. Defaults to True.
+        """
+        self._assert_key_column_names_match(snap_to_merge)
+        suffix = "_to_merge"
+        merged = pd.merge(
+            self.df,
+            snap_to_merge.df,
+            how="left",
+            on=self.file_name_col,
+            suffixes=(None, suffix),
+        )
+        updated_labels = []
+        unmatched_files = []
+        for _, row in merged.iterrows():
+            to_merge_label = f"{self.label_col}{suffix}"
+            labels = row[self.label_col]
+            if isinstance(row[to_merge_label], (dict, list)):
+                if ensure_identical_text:
+                    try:
+                        assert row[self.text_col] == row[f"{self.text_col}{suffix}"]
+                    except AssertionError:
+                        raise ToolkitInputError(
+                            f"Text from {row[self.file_name_col]} doesn't match"
+                        )
+                labels["targets"].extend(row[to_merge_label]["targets"])
+            else:
+                unmatched_files.append(row[self.file_name_col])
+            updated_labels.append(labels)
+        if unmatched_files:
+            print(f"No new labels were added for these files: {unmatched_files}")
+        merged[self.label_col] = updated_labels
+        self.df = merged.drop([i for i in merged.columns if i.endswith(suffix)], axis=1)
+
+    def to_csv(
+        self,
+        path: str,
+        only_keep_key_columns: bool = True,
+        convert_label_col_to_json: bool = True,
+    ):
+
+        if only_keep_key_columns:
+            self.drop_unneeded_columns()
+        if convert_label_col_to_json:
+            self._convert_col_to_json(self.label_col)
+        self.df.to_csv(path, index=False)
+
+    def get_all_labeled_text(
+        self, label_name: str, return_per_document: bool = False
+    ) -> Union[List[List[str]], List[str]]:
+        """
+        Get all of the text that was tagged for a given label
+        Args:
+            label_name (str): name of the label
+            return_per_document (bool, optional): return a list per document or one list
+                with everything. Defaults to False.
+        """
+        available_labels = self.get_extraction_label_names()
+        if label_name not in available_labels:
+            raise ToolkitInputError(
+                f"'{label_name}' not present among available labels: {available_labels}"
+            )
+        all_labeled_text = []
+        for text, labels in zip(self.df[self.text_col], self.df[self.label_col]):
+            text_found = []
+            for lab in labels["targets"]:
+                if lab["label"] == label_name:
+                    for span in lab["spans"]:
+                        text_found.append(text[span["start"] : span["end"]])
+            if return_per_document:
+                all_labeled_text.append(text_found)
+            else:
+                all_labeled_text.extend(text_found)
+        return all_labeled_text
+
+    def split_and_write_to_csv(
+        self, output_dir: str, num_splits: int = 5, output_base_name: str = "split_num"
+    ) -> None:
+        """
+        For large files that may face memory constraints, split the file into multiple
+        CSVs and write to disk.
+
+        Args:
+            output_dir (str): Location where split files will be written.
+            num_splits (int, optional): The number of splits of the CSV. Defaults to 5.
+            output_base_name (str, optional): The base name of the split CSVs:
+                Defaults to "split_num". So files would be "split_num_1.csv",
+                "split_num_2.csv", etc.
+        """
+        split_length = self.number_of_samples // num_splits
+        rows_taken = 0
+        for i in range(1, num_splits + 1):
+            if i == num_splits:
+                sub_df = self.df.iloc[rows_taken:]
+            else:
+                sub_df = self.df.iloc[rows_taken : rows_taken + split_length]
+            split_file_loc = os.path.join(output_dir, f"{output_base_name}_{i}.csv")
+            sub_df.to_csv(split_file_loc, index=False)
+            rows_taken += split_length
+            print(f"Wrote split {i} of {num_splits}: {split_file_loc}")
+
+    def update_label_col_format(self, task_type: str = "annotation"):
+        """
+        Modifies label column to updated format with spans.
+        Args:
+            task_type (str): Task type to specifiy df as
+                (annotation = extraction, classification = classification)
+
+        Note: page_num in the span is currently set to None.
+        """
+        updated_column = []
+        for label_set in self.df[self.label_col]:
+            updated_targets = []
+            for span in label_set:
+                updated_span = {
+                    "label": span["label"],
+                    "spans": [
+                        {"start": span["start"], "page_num": None, "end": span["end"]}
+                    ],
+                }
+                updated_targets.append(updated_span)
+
+            updated_column.append({"task_type": task_type, "targets": updated_targets})
+        self.df[self.label_col] = updated_column
+
+    def __eq__(self, other: "Snapshot"):
+        """
+        Check if two snapshots can be merged based on common column names
+        """
+        return (
+            self.label_col == other.label_col
+            and self.text_col == other.text_col
+            and self.file_name_col == other.file_name_col
+        )
+
+    def _convert_col_from_json(self):
+        try:
+            self.df[self.label_col] = self.df[self.label_col].apply(json.loads)
+        except (TypeError, JSONDecodeError):
+            if isinstance(self.df[self.label_col].iloc[0], list):
+                return  # json column already converted
+            raise ToolkitInputError(
+                f"{self.label_col} doesn't contain valid extraction labels"
+            )
+
+    def _assert_key_column_names_match(self, snapshot: "Snapshot"):
+        try:
+            assert self == snapshot
+        except AssertionError:
+            raise ToolkitInputError(
+                f"Column names don't match for this Snapshot: {snapshot}"
+            )
+
+    def _convert_col_to_json(self, column: str):
+        self.df[column] = self.df[column].apply(json.dumps)
+
+    def _infer_text_col(self):
+        if "text" in self.df.columns:
+            self.text_col = "text"
+        elif "document" in self.df.columns:
+            self.text_col = "document"
+        elif "source" in self.df.columns:
+            self.text_col = "source"
+        else:
+            error_message = self._format_column_error_message("text_col")
+            raise ToolkitInstantiationError(error_message)
+
+    def _infer_label_col(self):
+        question_col = [
+            col for col in self.df.columns if "question" in col or "target" in col
+        ]
+        if len(question_col) != 1:
+            error_message = self._format_column_error_message("label_col")
+            raise ToolkitInstantiationError(error_message)
+        self.label_col = question_col[0]
+
+    def _infer_file_name_col(self):
+        file_name_col = [col for col in self.df.columns if "file_name" in col]
+        if len(file_name_col) != 1:
+            error_message = self._format_column_error_message("file_name_col")
+            raise ToolkitInstantiationError(error_message)
+        self.file_name_col = file_name_col[0]
+
+    def _format_column_error_message(self, column):
+        column_list = list(self.df.columns)
+        return f"You must set attribute '{column}' from {column_list}"
+
+    def __repr__(self):
+        return f"Snapshot, label_col: {self.label_col}, loc: {self.path_to_snapshot}"
+
+    @property
+    def number_of_samples(self) -> int:
+        return self.df.shape[0]
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        return self.df.shape

indico_toolkit/structure/__init__.py

@@ -0,0 +1,3 @@
+from .create_structure import Structure
+
+__all__ = ("Structure",)

indico_toolkit/structure/create_structure.py

@@ -0,0 +1,232 @@
+import json
+import os
+import shutil
+import tempfile
+from typing import List
+
+from indico.queries import (
+    AddModelGroupComponent,
+    CreateDataset,
+    CreateWorkflow,
+    GetDataset,
+    GetWorkflow,
+    NewLabelsetArguments,
+)
+from indico.types import (
+    Dataset,
+    OcrEngine,
+    TableReadOrder,
+    Workflow,
+)
+
+from ..errors import ToolkitInputError
+from .queries import GetExampleIds, GetTeachDetails, LabelTeachTask
+from .utils import ModelTaskType
+
+
+class Structure:
+    def __init__(self, client):
+        self.client = client
+
+    def create_dataset(
+        self,
+        dataset_name: str,
+        files_to_upload: List[str],
+        read_api: bool = True,
+        **kwargs,
+    ) -> Dataset:
+        """
+        Creates a dataset from a list of files.
+
+        Args:
+            name_of_dataset (str): Name of the created dataset
+            file_path (str): Path of the file to copy.
+            read_api (bool, optional): OCR Engine used for the dataset.
+                Defaults to True=READ_API / False=OMNIPAGE
+        Kwargs:
+            Advanced OCR settings
+        """
+        read_api_settings = {
+            "auto_rotate": True,
+            "single_column": False,
+            "upscale_images": True,
+            "languages": ["ENG"],
+        }
+        omnipage_settings = {
+            **read_api_settings,
+            "force_render": True,
+            "native_layout": False,
+            "native_pdf": False,
+            "table_read_order": "row",
+        }
+        for arg in kwargs.keys():
+            if read_api:
+                read_api_settings.update({arg: kwargs[arg]})
+            else:
+                if arg == "table_read_order" and kwargs[arg] not in ["row", "column"]:
+                    raise ToolkitInputError(
+                        f"Keyword argument {arg} got an unexpected value of "
+                        f"{kwargs[arg]}, expected value of either 'row' or 'column'"
+                    )
+                omnipage_settings.update({arg: kwargs[arg]})
+
+        ocr_engine = OcrEngine.READAPI if read_api else OcrEngine.OMNIPAGE
+        if read_api:
+            ocr_settings = {"read_api_ocr_options": read_api_settings}
+        else:
+            omnipage_settings["table_read_order"] = (
+                TableReadOrder.ROW
+                if omnipage_settings["table_read_order"] == "row"
+                else TableReadOrder.COLUMN
+            )
+            ocr_settings = {"omnipage_ocr_options": omnipage_settings}
+
+        dataset = self.client.call(
+            CreateDataset(
+                dataset_name,
+                files=files_to_upload,
+                dataset_type="DOCUMENT",
+                ocr_engine=ocr_engine,
+                **ocr_settings,
+            )
+        )
+        print(f"Dataset created with ID: {dataset.id}")
+        return dataset
+
+    def create_duplicate_dataset(
+        self,
+        dataset_name: str,
+        file_path: str,
+        times_to_copy_files=55,
+        read_api=True,
+        **kwargs,
+    ) -> Dataset:
+        """
+        Creates a dataset w/ duplicate instances of 1 file, historically used to create
+        a spoofed demo.
+
+        Args:
+            file_path (str): Path of the file to copy.
+            name_of_dataset (str): Name of the created dataset
+            times_to_copy_files (int, optional): Amount of times to copy the file.
+                Defaults to 55.
+            read_api (bool, optional): OCR Engine used for the dataset.
+                Defaults to True=READ_API / False=OMNIPAGE
+        Kwargs:
+            Advanced OCR settings
+        """
+        tempdir = tempfile.TemporaryDirectory()
+        try:
+            files_to_add = []
+            for i in range(times_to_copy_files):
+                fpath = os.path.join(tempdir.name, f"sample_{i+1}.pdf")
+                shutil.copy(file_path, fpath)
+                files_to_add.append(fpath)
+            self.create_dataset(
+                dataset_name=dataset_name,
+                files_to_upload=files_to_add,
+                read_api=read_api,
+                **kwargs,
+            )
+        except Exception as e:
+            print(e)
+        finally:
+            tempdir.cleanup()
+
+    def create_workflow(self, name: str, dataset_id: int) -> Workflow:
+        """
+
+        Args:
+            name (str): Name of the workflow to be created
+            dataset_id (int): Dataset ID for the newly created workflow
+        """
+        workflow = self.client.call(CreateWorkflow(name=name, dataset_id=dataset_id))
+        print(f"Workflow created with ID: {workflow.id}")
+        return workflow
+
+    def add_teach_task(
+        self,
+        task_name: str,
+        labelset_name: str,
+        target_names: List[str],
+        dataset_id: int,
+        workflow_id: int,
+        model_type: str = "annotation",
+        data_column: str = "document",
+        prev_comp_id: int = None,
+        **kwargs,
+    ) -> Workflow:
+        """
+        Args:
+            task_name (str): Teach task name
+            labelset_name (str): Name for created labelset
+            target_names (List[str]): List of target (label) names
+            dataset_id (int): Dataset ID.
+            workflow_id (int): Workflow ID.
+            model_type (str, optional): Defaults to annotation.
+                Available model types:
+                    classification
+                    annotation
+                    classification_unbundling
+            data_column (str, optional): Defaults to "document".
+        Kwargs:
+            Advanced model training options
+        Returns:
+            Updated Workflow with newly added model group component.
+        """
+        workflow = self.client.call(GetWorkflow(workflow_id))
+        dataset = self.client.call(GetDataset(dataset_id))
+        model_map = {
+            "classification": ModelTaskType.CLASSIFICATION,
+            "annotation": ModelTaskType.ANNOTATION,
+            "classification_unbundling": ModelTaskType.CLASSIFICATION_UNBUNDLING,
+            "classification_multiple": ModelTaskType.CLASSIFICATION_MULTIPLE,
+        }
+        if model_type not in model_map.keys():
+            raise ToolkitInputError(
+                f"{model_type} not found. Available options "
+                f"include {[model for model in model_map.keys()]}"
+            )
+        workflow = self.client.call(GetWorkflow(workflow_id))
+        if not prev_comp_id:
+            prev_comp_id = workflow.component_by_type("INPUT_OCR_EXTRACTION").id
+
+        column_id = dataset.datacolumn_by_name(data_column).id
+        new_labelset = NewLabelsetArguments(
+            labelset_name,
+            datacolumn_id=column_id,
+            task_type=model_map[model_type],
+            target_names=target_names,
+        )
+
+        workflow = self.client.call(
+            AddModelGroupComponent(
+                workflow_id=workflow.id,
+                name=task_name,
+                dataset_id=dataset.id,
+                source_column_id=column_id,
+                after_component_id=prev_comp_id,
+                new_labelset_args=new_labelset,
+                model_training_options=json.dumps(kwargs),
+            )
+        )
+        print(
+            "Newly created teach task with teach_id: "
+            f"{workflow.components[-1].model_group.questionnaire_id}"
+        )
+        return workflow
+
+    def get_teach_details(self, teach_task_id: int):
+        return self.client.call(GetTeachDetails(teach_task_id=teach_task_id))
+
+    def get_example_ids(self, model_group_id: int, limit: int):
+        return self.client.call(
+            GetExampleIds(model_group_id=model_group_id, limit=limit)
+        )
+
+    def label_teach_task(self, label_set_id: int, labels: dict, model_group_id: int):
+        return self.client.call(
+            LabelTeachTask(
+                label_set_id=label_set_id, labels=labels, model_group_id=model_group_id
+            )
+        )

indico_toolkit/structure/queries.py

@@ -0,0 +1,119 @@
+from indico.queries import GraphQLRequest
+
+
+class GetTeachDetails(GraphQLRequest):
+    GET_TEACH_DETAILS = """
+        query getCrowdlabelQuestionnaire($teach_task_id: Int!) {
+            questionnaire(id: $teach_task_id) {
+                id
+                question {
+                    id
+                    type
+                    text
+                    modelGroupId
+                    status
+                    labelset {
+                        id
+                        targetNames {
+                            id
+                            name
+                            position
+                        }
+                    }
+                }
+            }
+        }
+    """
+
+    def __init__(self, *, teach_task_id: int):
+        super().__init__(
+            self.GET_TEACH_DETAILS,
+            variables={"teach_task_id": teach_task_id},
+        )
+
+    def process_response(self, response):
+        return super().process_response(response)
+
+
+class GetExampleIds(GraphQLRequest):
+    GET_EXAMPLES = """
+        query getExamplesList(
+            $modelGroupId: Int!,
+            $filters: ExampleFilter,
+            $skip: Int,
+            $before: Int,
+            $after: Int,
+            $limit: Int,
+            $desc: Boolean,
+            $orderBy: ExampleOrder
+        ) {
+          modelGroup(modelGroupId: $modelGroupId) {
+            id
+            pagedExamples(
+                filters: $filters,
+                skip: $skip,
+                before: $before,
+                after: $after,
+                limit: $limit,
+                desc: $desc,
+                orderBy:
+                $orderBy
+            ) {
+              examples {
+                id
+                datarowId
+                updatedAt
+                status
+                datafile {
+                    name
+                }
+              }
+              pageInfo {
+                startCursor
+                endCursor
+                hasNextPage
+                aggregateCount
+              }
+            }
+          }
+        }
+    """
+
+    def __init__(self, *, model_group_id: int, limit: int):
+        super().__init__(
+            self.GET_EXAMPLES,
+            variables={"modelGroupId": model_group_id, "limit": limit},
+        )
+
+    def process_response(self, response):
+        return super().process_response(response)
+
+
+class LabelTeachTask(GraphQLRequest):
+    LABEL_TASK = """
+        mutation submitQuestionnaireExample
+            ($labelsetId: Int!,
+            $labels: [LabelInput]!, 
+            $modelGroupId: Int) {
+          submitLabelsV2(
+          labels: $labels, 
+          labelsetId: $labelsetId, 
+          modelGroupId: $modelGroupId) {
+            success
+            __typename
+          }
+        }
+    """
+
+    def __init__(self, *, label_set_id: int, labels, model_group_id: int):
+        super().__init__(
+            self.LABEL_TASK,
+            variables={
+                "labelsetId": label_set_id,
+                "labels": labels,
+                "modelGroupId": model_group_id,
+            },
+        )
+
+    def process_response(self, response):
+        return super().process_response(response)