Add settings for model averaging in the frontend (#620)

Dref360 · web-flow · commit 3102ddc17c55 · 2024-01-22T17:03:58.000Z
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ Released changes are shown in the
 ## [Not released]
 
 ### Added
+* Prediction after BMA can now be displayed in the app.
 
 ### Changed
 
diff --git a/azimuth/modules/base_classes/aggregation_module.py b/azimuth/modules/base_classes/aggregation_module.py
@@ -31,7 +31,7 @@ class FilterableModule(AggregationModule[ConfigScope], ExpirableMixin, ABC):
     """Filterable Module are affected by filters in mod options."""
 
     required_mod_options = {"pipeline_index"}
-    optional_mod_options = {"filters", "without_postprocessing"}
+    optional_mod_options = {"filters", "without_postprocessing", "use_bma"}
 
     def get_dataset_split(self, name: DatasetSplitName = None) -> Dataset:
         """Get the specified dataset_split, filtered according to mod_options.
diff --git a/azimuth/modules/model_contracts/text_classification.py b/azimuth/modules/model_contracts/text_classification.py
@@ -109,6 +109,7 @@ def post_process(self, batch: Dataset) -> List[PredictionResponse]:
             mod_options=ModuleOptions(
                 model_contract_method_name=SupportedMethod.Predictions,
                 pipeline_index=self.mod_options.pipeline_index,
+                use_bma=self.mod_options.use_bma,
                 indices=cast(List[int], batch[DatasetColumn.row_idx]),
             ),
         )
diff --git a/azimuth/modules/model_performance/confidence_binning.py b/azimuth/modules/model_performance/confidence_binning.py
@@ -114,7 +114,7 @@ class ConfidenceBinIndexModule(DatasetResultModule[ModelContractConfig]):
     """Return confidence bin indices for the selected dataset split."""
 
     required_mod_options = {"pipeline_index"}
-    optional_mod_options = DatasetResultModule.optional_mod_options | {"threshold"}
+    optional_mod_options = DatasetResultModule.optional_mod_options | {"threshold", "use_bma"}
 
     def compute_on_dataset_split(self) -> List[int]:  # type: ignore
         """Get the bin indices for each utterance in the dataset split.
diff --git a/azimuth/modules/model_performance/outcome_count.py b/azimuth/modules/model_performance/outcome_count.py
@@ -178,7 +178,7 @@ class OutcomeCountPerThresholdModule(AggregationModule[ModelContractConfig]):
     """Compute the outcome count per threshold."""
 
     required_mod_options = {"pipeline_index"}
-    optional_mod_options = {"x_ticks_count"}
+    optional_mod_options = {"x_ticks_count", "use_bma"}
 
     def compute_on_dataset_split(self) -> List[OutcomeCountPerThresholdResponse]:  # type: ignore
         if not postprocessing_editable(self.config, self.mod_options.pipeline_index):
@@ -199,6 +199,7 @@ def compute_on_dataset_split(self) -> List[OutcomeCountPerThresholdResponse]:  #
                     # Convert to float instead of numpy.float64
                     threshold=float(th),
                     pipeline_index=self.mod_options.pipeline_index,
+                    use_bma=self.mod_options.use_bma,
                 ),
             )
             outcomes = outcomes_mod.compute_on_dataset_split()
diff --git a/azimuth/modules/model_performance/outcomes.py b/azimuth/modules/model_performance/outcomes.py
@@ -20,7 +20,7 @@ class OutcomesModule(DatasetResultModule[ModelContractConfig]):
     """Computes the outcome for each utterance in the dataset split."""
 
     required_mod_options = {"pipeline_index"}
-    optional_mod_options = DatasetResultModule.optional_mod_options | {"threshold"}
+    optional_mod_options = DatasetResultModule.optional_mod_options | {"threshold", "use_bma"}
 
     def _get_predictions(self, without_postprocessing: bool) -> ndarray:
         mod_options = self.mod_options.copy(deep=True)
diff --git a/azimuth/routers/app.py b/azimuth/routers/app.py
@@ -104,6 +104,8 @@ def get_dataset_info(
             eval=eval_dm and eval_dm.num_rows, train=training_dm and training_dm.num_rows
         ),
         similarity_available=similarity_available(config),
+        model_averaging_available=config.uncertainty is not None
+        and config.uncertainty.iterations > 1,
         postprocessing_editable=None
         if config.pipelines is None
         else [postprocessing_editable(config, idx) for idx in range(len(config.pipelines))],
diff --git a/azimuth/routers/model_performance/confidence_histogram.py b/azimuth/routers/model_performance/confidence_histogram.py
@@ -31,10 +31,12 @@ def get_confidence_histogram(
     dataset_split_manager: DatasetSplitManager = Depends(get_dataset_split_manager),
     pipeline_index: int = Depends(require_pipeline_index),
     without_postprocessing: bool = Query(False, title="Without Postprocessing"),
+    use_bma: bool = Query(False, title="Use BMA"),
 ) -> ConfidenceHistogramResponse:
     mod_options = ModuleOptions(
         filters=named_filters.to_dataset_filters(dataset_split_manager.get_class_names()),
         pipeline_index=pipeline_index,
+        use_bma=use_bma,
         without_postprocessing=without_postprocessing,
     )
 
diff --git a/azimuth/routers/model_performance/confusion_matrix.py b/azimuth/routers/model_performance/confusion_matrix.py
@@ -31,12 +31,14 @@ def get_confusion_matrix(
     dataset_split_manager: DatasetSplitManager = Depends(get_dataset_split_manager),
     pipeline_index: int = Depends(require_pipeline_index),
     without_postprocessing: bool = Query(False, title="Without Postprocessing"),
+    use_bma: bool = Query(False, title="Use BMA"),
     normalize: bool = Query(True, title="Normalize"),
     reorder_classes: bool = Query(True, title="Reorder Classes"),
 ) -> ConfusionMatrixResponse:
     mod_options = ModuleOptions(
         filters=named_filters.to_dataset_filters(dataset_split_manager.get_class_names()),
         pipeline_index=pipeline_index,
+        use_bma=use_bma,
         without_postprocessing=without_postprocessing,
         cf_normalize=normalize,
         cf_reorder_classes=reorder_classes,
diff --git a/azimuth/routers/model_performance/metrics.py b/azimuth/routers/model_performance/metrics.py
@@ -39,10 +39,12 @@ def get_metrics(
     dataset_split_manager: DatasetSplitManager = Depends(get_dataset_split_manager),
     pipeline_index: int = Depends(require_pipeline_index),
     without_postprocessing: bool = Query(False, title="Without Postprocessing"),
+    use_bma: bool = Query(False, title="Use BMA"),
 ) -> MetricsAPIResponse:
     mod_options = ModuleOptions(
         filters=named_filters.to_dataset_filters(dataset_split_manager.get_class_names()),
         pipeline_index=pipeline_index,
+        use_bma=use_bma,
         without_postprocessing=without_postprocessing,
     )
 
diff --git a/azimuth/routers/model_performance/outcome_count.py b/azimuth/routers/model_performance/outcome_count.py
@@ -35,10 +35,9 @@ def get_outcome_count_per_threshold(
     task_manager: TaskManager = Depends(get_task_manager),
     dataset_split_manager: DatasetSplitManager = Depends(get_dataset_split_manager),
     pipeline_index: int = Depends(require_pipeline_index),
+    use_bma: bool = Query(False, title="Use BMA"),
 ) -> OutcomeCountPerThresholdResponse:
-    mod_options = ModuleOptions(
-        pipeline_index=pipeline_index,
-    )
+    mod_options = ModuleOptions(pipeline_index=pipeline_index, use_bma=use_bma)
 
     try:
         task_result: List[OutcomeCountPerThresholdResponse] = get_standard_task_result(
diff --git a/azimuth/startup.py b/azimuth/startup.py
@@ -65,7 +65,21 @@ class Startup:
         SupportedMethod.Predictions,
         mod_options={"use_bma": True},
         run_on_all_pipelines=True,
-    )
+    ),
+    Startup(
+        "outcomes_bma",
+        SupportedModule.Outcome,
+        dependency_names=["prediction_bma"],
+        mod_options={"use_bma": True},
+        run_on_all_pipelines=True,
+    ),
+    Startup(
+        "confidence_bins_bma",
+        SupportedModule.ConfidenceBinIndex,
+        dependency_names=["prediction_bma"],
+        mod_options={"use_bma": True},
+        run_on_all_pipelines=True,
+    ),
 ]
 
 PERTURBATION_TESTING_TASKS = [
diff --git a/azimuth/types/app.py b/azimuth/types/app.py
@@ -43,6 +43,7 @@ class DatasetInfoResponse(AliasModel):
         ..., title="Utterance count per dataset split."
     )
     similarity_available: bool = Field(..., title="Indicator if similarities are available.")
+    model_averaging_available: bool = Field(..., title="Indicator if model averaging is available.")
     postprocessing_editable: Optional[List[bool]] = Field(
         ..., title="Indicator if postprocessing can be overwritten.", nullable=True
     )
diff --git a/tests/test_routers/test_app.py b/tests/test_routers/test_app.py
@@ -52,6 +52,7 @@ def test_get_dataset_info(app: FastAPI) -> None:
         "dataActions": ALL_DATA_ACTION_FILTERS,
         "similarityAvailable": True,
         "smartTags": ALL_SMART_TAG_FILTERS,
+        "modelAveragingAvailable": False,
     }
 
 
diff --git a/webapp/src/components/Analysis/UtterancesTable.tsx b/webapp/src/components/Analysis/UtterancesTable.tsx
@@ -62,6 +62,7 @@ import {
   QueryPaginationState,
   QueryPipelineState,
   QueryPostprocessingState,
+  QueryBMAState,
   QueryState,
 } from "types/models";
 import {
@@ -145,6 +146,7 @@ type Props = {
   pagination: QueryPaginationState;
   pipeline: QueryPipelineState;
   postprocessing: QueryPostprocessingState;
+  modelAveraging: QueryBMAState;
 };
 
 const UtterancesTable: React.FC<Props> = ({
@@ -157,6 +159,7 @@ const UtterancesTable: React.FC<Props> = ({
   pagination,
   pipeline,
   postprocessing,
+  modelAveraging,
 }) => {
   const history = useHistory();
   const classes = useStyles();
@@ -175,6 +178,7 @@ const UtterancesTable: React.FC<Props> = ({
     offset,
     ...pipeline,
     ...postprocessing,
+    ...modelAveraging,
   };
 
   const {
@@ -212,6 +216,7 @@ const UtterancesTable: React.FC<Props> = ({
         ...pagination,
         ...pipeline,
         ...postprocessing,
+        ...modelAveraging,
         ...update,
       }
     )}`;
diff --git a/webapp/src/components/ConfidenceHistogramTopWords.tsx b/webapp/src/components/ConfidenceHistogramTopWords.tsx
@@ -15,6 +15,7 @@ import {
   QueryPaginationState,
   QueryPipelineState,
   QueryPostprocessingState,
+  QueryBMAState,
   WordCount,
 } from "types/models";
 import { ALL_OUTCOMES } from "utils/const";
@@ -26,6 +27,7 @@ type Props = {
   pagination: QueryPaginationState;
   pipeline: Required<QueryPipelineState>;
   postprocessing: QueryPostprocessingState;
+  modelAveraging: QueryBMAState;
 };
 
 const ConfidenceHistogramTopWords: React.FC<Props> = ({
@@ -35,6 +37,7 @@ const ConfidenceHistogramTopWords: React.FC<Props> = ({
   pagination,
   pipeline,
   postprocessing,
+  modelAveraging,
 }) => {
   const { jobId, datasetSplitName } = useParams<{
     jobId: string;
@@ -58,6 +61,7 @@ const ConfidenceHistogramTopWords: React.FC<Props> = ({
     ...filtersWithoutBins,
     ...pipeline,
     ...postprocessing,
+    ...modelAveraging,
   });
 
   const { data: topWords, isFetching: isFetchingTopWords } =
@@ -67,6 +71,7 @@ const ConfidenceHistogramTopWords: React.FC<Props> = ({
       ...filters,
       ...pipeline,
       ...postprocessing,
+      ...modelAveraging,
     });
 
   const topWordsCounts = topWords;
diff --git a/webapp/src/components/ConfusionMatrix.tsx b/webapp/src/components/ConfusionMatrix.tsx
@@ -19,6 +19,7 @@ import {
   QueryFilterState,
   QueryPipelineState,
   QueryPostprocessingState,
+  QueryBMAState,
 } from "types/models";
 import { OUTCOME_COLOR, UNKNOWN_ERROR } from "utils/const";
 import { classNames, constructSearchString } from "utils/helpers";
@@ -36,6 +37,7 @@ type Props = {
   filters: QueryFilterState;
   pipeline: Required<QueryPipelineState>;
   postprocessing: QueryPostprocessingState;
+  modelAveraging: QueryBMAState;
   predictionFilters?: string[];
   labelFilters?: string[];
 };
@@ -93,6 +95,7 @@ const ConfusionMatrix: React.FC<Props> = ({
   filters,
   pipeline,
   postprocessing,
+  modelAveraging,
   predictionFilters,
   labelFilters,
 }) => {
@@ -114,6 +117,7 @@ const ConfusionMatrix: React.FC<Props> = ({
       ...filters,
       ...pipeline,
       ...postprocessing,
+      ...modelAveraging,
     });
 
   if (isLoading) {
@@ -140,6 +144,7 @@ const ConfusionMatrix: React.FC<Props> = ({
         ...filters,
         ...pipeline,
         ...postprocessing,
+        ...modelAveraging,
         ...newConfusionMatrix,
       })}`
     );
@@ -152,6 +157,7 @@ const ConfusionMatrix: React.FC<Props> = ({
         ...filters,
         ...pipeline,
         ...postprocessing,
+        ...modelAveraging,
         prediction: [data.classNames[columnIndex]],
         label: [data.classNames[rowIndex]],
       })}`}
@@ -299,6 +305,7 @@ const ConfusionMatrix: React.FC<Props> = ({
                 ...filters,
                 ...pipeline,
                 ...postprocessing,
+                ...modelAveraging,
                 prediction: [data.classNames[i]],
               })}`}
               key={`column-${i}`}
@@ -333,6 +340,7 @@ const ConfusionMatrix: React.FC<Props> = ({
                 ...filters,
                 ...pipeline,
                 ...postprocessing,
+                ...modelAveraging,
                 label: [data.classNames[i]],
               })}`}
               key={`row-${i}`}
diff --git a/webapp/src/components/Controls/Controls.tsx b/webapp/src/components/Controls/Controls.tsx
diff --git a/webapp/src/components/Metrics/Metrics.tsx b/webapp/src/components/Metrics/Metrics.tsx
diff --git a/webapp/src/components/PageHeader.tsx b/webapp/src/components/PageHeader.tsx
diff --git a/webapp/src/mocks/api/mockDatasetInfoAPI.ts b/webapp/src/mocks/api/mockDatasetInfoAPI.ts
diff --git a/webapp/src/pages/Exploration.tsx b/webapp/src/pages/Exploration.tsx
diff --git a/webapp/src/types/generated/generatedTypes.ts b/webapp/src/types/generated/generatedTypes.ts
diff --git a/webapp/src/types/models.ts b/webapp/src/types/models.ts
diff --git a/webapp/src/utils/helpers.ts b/webapp/src/utils/helpers.ts

Original file line number	Diff line number	Diff line change
`@@ -109,6 +109,7 @@ def post_process(self, batch: Dataset) -> List[PredictionResponse]:`
`109`	`109`	`mod_options=ModuleOptions(`
`110`	`110`	`model_contract_method_name=SupportedMethod.Predictions,`
`111`	`111`	`pipeline_index=self.mod_options.pipeline_index,`
	`112`	`+ use_bma=self.mod_options.use_bma,`
`112`	`113`	`indices=cast(List[int], batch[DatasetColumn.row_idx]),`
`113`	`114`	`),`
`114`	`115`	`)`