Handle perturbations in run predictions frontend

helm-frontend/src/components/InstanceData.tsx

@@ -20,11 +20,14 @@ export default function InstanceData({
   predictions,
   metricFieldMap,
 }: Props) {
+  const renderInstanceId = (instance: Instance): string => {
+    return instance.perturbation === undefined
+      ? `Instance id: ${instance.id} [split: ${instance.split}]`
+      : `Instance id: ${instance.id} [split: ${instance.split}][perturbation: ${instance.perturbation.name}]`;
+  };
   return (
     <div className="border p-4">
-      <h3 className="text-xl mb-4">
-        {`Instance id: ${instance.id} [split:  ${instance.split}]`}
-      </h3>
+      <h3 className="text-xl mb-4">{renderInstanceId(instance)}</h3>
       <h3>Input</h3>
       {instance.input.multimedia_content !== undefined ? (
         <MultimediaObjectDisplay

helm-frontend/src/components/Predictions.tsx

@@ -30,7 +30,9 @@ export default function Predictions({
       <div className="flex flex-wrap justify-start items-start">
         {predictions.map((prediction, idx) => (
           <div className="w-full" key={idx}>
-            {predictions.length > 1 ? <h2>Trial {idx}</h2> : null}
+            {predictions.length > 1 ? (
+              <h2>Trial {prediction.train_trial_index}</h2>
+            ) : null}
             <div className="mt-2 w-full">
               {prediction.base64_images &&
               prediction.base64_images.length > 0 ? (

helm-frontend/src/routes/Run.tsx

@@ -13,9 +13,7 @@ import type Instance from "@/types/Instance";
 import getStatsByName from "@/services/getStatsByName";
 import type Stat from "@/types/Stat";
 import getDisplayRequestsByName from "@/services/getDisplayRequestsByName";
-import type DisplayRequestsMap from "@/types/DisplayRequestsMap";
 import getDisplayPredictionsByName from "@/services/getDisplayPredictionsByName";
-import type DisplayPredictionsMap from "@/types/DisplayPredictionsMap";
 import getScenarioByName from "@/services/getScenarioByName";
 import type Scenario from "@/types/Scenario";
 import type AdapterFieldMap from "@/types/AdapterFieldMap";
@@ -32,6 +30,8 @@ import MarkdownValue from "@/components/MarkdownValue";
 import StatNameDisplay from "@/components/StatNameDisplay";
 import getRunsToRunSuites from "@/services/getRunsToRunSuites";
 import getSuiteForRun from "@/services/getSuiteForRun";
+import DisplayPrediction from "@/types/DisplayPrediction";
+import DisplayRequest from "@/types/DisplayRequest";
 
 const INSTANCES_PAGE_SIZE = 10;
 const METRICS_PAGE_SIZE = 50;
@@ -45,10 +45,10 @@ export default function Run() {
   const [instances, setInstances] = useState<Instance[]>([]);
   const [stats, setStats] = useState<Stat[]>([]);
   const [displayPredictionsMap, setDisplayPredictionsMap] = useState<
-    DisplayPredictionsMap | undefined
+    undefined | { [key: string]: { [key: string]: DisplayPrediction[] } }
   >();
   const [displayRequestsMap, setDisplayRequestsMap] = useState<
-    DisplayRequestsMap | undefined
+    undefined | { [key: string]: { [key: string]: DisplayRequest[] } }
   >();
   const [currentInstancesPage, setCurrentInstancesPage] = useState<number>(1);
   const [totalInstancesPages, setTotalInstancesPages] = useState<number>(1);
@@ -81,13 +81,15 @@ export default function Run() {
         scenario,
         displayPredictions,
         displayRequests,
+        schema,
       ] = await Promise.all([
         getRunSpecs(signal),
         getInstances(runName, signal, suite),
         getStatsByName(runName, signal, suite),
         getScenarioByName(runName, signal, suite),
         getDisplayPredictionsByName(runName, signal, suite),
         getDisplayRequestsByName(runName, signal, suite),
+        getSchema(signal),
       ]);
 
       setRunSpec(runSpecs.find((rs) => rs.name === runName));
@@ -110,25 +112,42 @@ export default function Run() {
       setCurrentMetricsPage(
         Math.max(Math.min(metricPage, totalMetricsPages), 1),
       );
-      setDisplayPredictionsMap(
-        displayPredictions.reduce((acc, cur) => {
-          if (acc[cur.instance_id] === undefined) {
-            acc[cur.instance_id] = [];
-          }
-          acc[cur.instance_id].push(cur);
-          return acc;
-        }, {} as DisplayPredictionsMap),
-      );
-      setDisplayRequestsMap(
-        displayRequests.reduce((acc, cur) => {
-          if (acc[cur.instance_id] === undefined) {
-            acc[cur.instance_id] = [];
-          }
-          acc[cur.instance_id].push(cur);
-          return acc;
-        }, {} as DisplayRequestsMap),
-      );
-      const schema = await getSchema(signal);
+
+      const displayRequestsArgh: {
+        [key: string]: { [key: string]: DisplayRequest[] };
+      } = {};
+      displayRequests.forEach((displayRequest) => {
+        const instanceId = displayRequest.instance_id;
+        const perturbationName = displayRequest.perturbation?.name || "";
+        if (displayRequestsArgh[instanceId] === undefined) {
+          displayRequestsArgh[instanceId] = {};
+        }
+        if (displayRequestsArgh[instanceId][perturbationName] === undefined) {
+          displayRequestsArgh[instanceId][perturbationName] = [];
+        }
+        displayRequestsArgh[instanceId][perturbationName].push(displayRequest);
+      });
+      setDisplayRequestsMap(displayRequestsArgh);
+
+      const displayPredictionsArgh: {
+        [key: string]: { [key: string]: DisplayPrediction[] };
+      } = {};
+      displayPredictions.forEach((displayPrediction) => {
+        const instanceId = displayPrediction.instance_id;
+        const perturbationName = displayPrediction.perturbation?.name || "";
+        if (displayPredictionsArgh[instanceId] === undefined) {
+          displayPredictionsArgh[instanceId] = {};
+        }
+        if (
+          displayPredictionsArgh[instanceId][perturbationName] === undefined
+        ) {
+          displayPredictionsArgh[instanceId][perturbationName] = [];
+        }
+        displayPredictionsArgh[instanceId][perturbationName].push(
+          displayPrediction,
+        );
+      });
+      setDisplayPredictionsMap(displayPredictionsArgh);
 
       setMetricFieldMap(
         schema.metrics.reduce((acc, cur) => {
@@ -268,8 +287,16 @@ export default function Run() {
               <InstanceData
                 key={`${instance.id}-${idx}`}
                 instance={instance}
-                requests={displayRequestsMap[instance.id]}
-                predictions={displayPredictionsMap[instance.id]}
+                requests={
+                  displayRequestsMap[instance.id][
+                    instance.perturbation?.name || ""
+                  ]
+                }
+                predictions={
+                  displayPredictionsMap[instance.id][
+                    instance.perturbation?.name || ""
+                  ]
+                }
                 metricFieldMap={metricFieldMap}
               />
             ))}

helm-frontend/src/types/DisplayPrediction.ts

@@ -1,4 +1,5 @@
 import CompletionAnnotation from "./CompletionAnnotation";
+import Perturbation from "./Perturbation";
 
 export default interface DisplayPrediction {
   instance_id: string;
@@ -16,4 +17,5 @@ export default interface DisplayPrediction {
   base64_images?: string[] | undefined;
   // beware you will have to update this for future custom annotations
   annotations?: Record<string, Array<CompletionAnnotation>> | undefined;
+  perturbation?: Perturbation | undefined;
 }

helm-frontend/src/types/DisplayRequest.ts

@@ -1,4 +1,5 @@
 import MultimediaObject from "@/types/MultimediaObject";
+import Perturbation from "./Perturbation";
 
 export default interface DisplayRequest {
   instance_id: string;
@@ -18,4 +19,5 @@ export default interface DisplayRequest {
     top_p: number;
     multimodal_prompt: MultimediaObject | undefined;
   };
+  perturbation?: Perturbation | undefined;
 }

helm-frontend/src/types/Instance.ts

@@ -1,5 +1,6 @@
 import MultimediaObject from "@/types/MultimediaObject";
 import type Reference from "@/types/Reference";
+import Perturbation from "./Perturbation";
 
 export default interface Instance {
   id: string;
@@ -9,4 +10,5 @@ export default interface Instance {
     multimedia_content: MultimediaObject | undefined;
   };
   references: Reference[];
+  perturbation?: Perturbation | undefined;
 }