Merge pull request #159 from IndicoDataSolutions/unbundling_metrics

adding unbundling metrics support

indico_toolkit/metrics/metrics.py

@@ -9,6 +9,26 @@
 
 
 class ExtractionMetrics(IndicoWrapper):
+    """
+    Example usage:
+
+    metrics = ExtractionMetrics(client)
+    metrics.get_extraction_metrics(MODEL_GROUP_ID)
+
+    # get a pandas dataframe of field level results
+    df = metrics.get_metrics_df()
+    print(df.head())
+
+    # get metrics for a specific span type and/or model ID
+    df = metrics.get_metrics_df(span_type="exact", select_model_id=102)
+    print(df.head())
+
+    # write the results to a CSV (can also optionally pass span_type/model ID here as well)
+    metrics.to_csv("./my_metrics.pdf")
+
+    # get an interactive bar plot to visualize model improvement over time
+    metrics.bar_plot("./my_bar_plot.html")
+    """
     def __init__(self, client: IndicoClient):
         self.client = client
         self.raw_metrics: List[dict] = None
@@ -169,6 +189,112 @@ def to_csv(
         df.to_csv(output_path, index=False)
 
 
+
+class UnbundlingMetrics(ExtractionMetrics):
+    """
+    Example Usage:
+
+    um = UnbundlingMetrics(client)
+    um.get_metrics(1232)
+    um.line_plot("./my_metric_plot.html", metric="recall", title="Insurance Model Recall Improvement")
+
+    """
+    def get_metrics(self, model_group_id: int):
+        """
+        Collect all metrics available based on a Model Group ID for an Unbundling model
+        Args:
+            model_group_id (int): Model Group ID that you're interestd in (available within the Explain UI)
+        """
+        results = self.graphQL_request(METRIC_QUERY, {"modelGroupId": model_group_id})
+        if len(results["modelGroups"]["modelGroups"]) == 0:
+            raise ToolkitInputError(
+                f"There are no models associated with ID: {model_group_id}"
+            )
+        results = results["modelGroups"]["modelGroups"][0]["models"]
+        raw_metrics = []
+        included_models = []
+        labeled_samples = []
+        for r in results:
+            model_info = json.loads(r["modelInfo"])
+            if "total_number_of_examples" not in model_info or "metrics" not in model_info:
+                # some dictionaries don't come back with required fields... 
+                continue
+            labeled_samples.append(model_info["total_number_of_examples"])
+            included_models.append(r["id"])
+            raw_metrics.append(model_info["metrics"]["per_class_metrics"])
+        self.raw_metrics = raw_metrics
+        self.included_models = included_models
+        self.number_of_samples = {model_id:samples for model_id, samples in zip(included_models, labeled_samples)}
+
+
+    def get_metrics_df(self) -> pd.DataFrame:
+        cleaned_metrics = []
+        for model_id, metrics in zip(self.included_models, self.raw_metrics):
+            for class_name in metrics:
+                scores = metrics[class_name]["page"]
+                scores["field_name"] = class_name
+                scores["model_id"] = model_id
+                scores["number_of_samples"] = self.number_of_samples[model_id]
+                cleaned_metrics.append(scores)
+        df = pd.DataFrame(cleaned_metrics)
+        return df.sort_values(by=["field_name", "model_id"], ascending=False)
+
+
+    def line_plot(
+        self,
+        output_path: str,
+        metric: str = "f1_score",
+        plot_title: str = "",
+        ids_to_exclude: List[int] = [],
+        fields_to_exclude: List[str] = [],
+    ):
+        """
+        Write an html line plot to disc with # of samples on x-axis, a metric on the y-axis and
+        each line representing a distinct field.
+        Will also open the plot automatically in your browser, where you will interactive
+        functionality and the ability to download a copy as a PNG as well.
+
+        Args:
+            output_path (str): where you want to write plot, e.g. "./myplot.html"
+            span_type (str): options include 'superset', 'exact', 'overlap' or 'token'
+            metric (str, optional): possible values are 'precision', 'recall', 'f1_score', 'false_positives',
+                                    'false_negatives', 'true_positives'. Defaults to "f1_score".
+            plot_title (str, optional): Title of the plot. Defaults to "".
+            ids_to_exclude (List[int], optional): Model Ids to exclude from plot.
+            fields_to_exclude (List[str], optional): Field Names to exclude from plot.
+        """
+        df = self.get_metrics_df()
+        if ids_to_exclude:
+            df = df.drop(df.loc[df["model_id"].isin(ids_to_exclude)].index)
+        if fields_to_exclude:
+            df = df.drop(df.loc[df["field_name"].isin(fields_to_exclude)].index)
+        df = df.sort_values(by=["field_name", "number_of_samples", metric])
+        plotting = Plotting()
+        for field in sorted(df["field_name"].unique()):
+            sub_df = df.loc[df["field_name"] == field].copy()
+            # ensure only one value per # of samples
+            sub_df = sub_df.drop_duplicates(subset=["number_of_samples"])
+            plotting.add_line_data(
+                sub_df["number_of_samples"],
+                sub_df[metric],
+                name=field,
+                color=None,
+            )
+        plotting.define_layout(
+            xaxis_title="Number of Samples",
+            legend_title="Field",
+            plot_title=plot_title,
+            yaxis_title=metric,
+        )
+        plotting.plot(output_path)
+
+    def bar_plot(self):
+        raise NotImplementedError("Bar Plot is not currently implemented for unbundling")
+
+    def get_extraction_metrics(self, model_group_id: int):
+        raise NotImplementedError("Not available for unbundling class")
+
+
 METRIC_QUERY = """ 
 query modelGroupMetrics($modelGroupId: Int!){
             modelGroups(

tests/indico_wrapper/test_dataset.py

@@ -1,3 +1,6 @@
+"""
+Test Datasets class methods
+"""
 import pytest
 from indico_toolkit.indico_wrapper import Datasets
 from indico.types import Dataset