From 3b810e1118fbb510152a6e8f2fab5b8e74b4a257 Mon Sep 17 00:00:00 2001
From: Yifan Mai <yifan@cs.stanford.edu>
Date: Fri, 31 May 2024 16:52:56 -0700
Subject: [PATCH] Allow non-lists to be displayed in AnnotationsDisplay

---
 .../src/components/AnnotationsDisplay.tsx     | 70 +++++++++++++------
 1 file changed, 48 insertions(+), 22 deletions(-)

diff --git a/helm-frontend/src/components/AnnotationsDisplay.tsx b/helm-frontend/src/components/AnnotationsDisplay.tsx
index 3bd7fbe8a95..0d44b8bd802 100644
--- a/helm-frontend/src/components/AnnotationsDisplay.tsx
+++ b/helm-frontend/src/components/AnnotationsDisplay.tsx
@@ -2,12 +2,56 @@ import CompletionAnnotation from "@/types/CompletionAnnotation";
 import Preview from "./Preview";
 import MediaObjectDisplay from "./MediaObjectDisplay";
 
+// TODO: This is a dirty hack to support annotations from
+// Image2Structure and AIRBench, but eventually we should make sure
+// all annotations are supported generally.
 type Props = {
   predictionAnnotations:
-    | Record<string, Array<CompletionAnnotation>>
+    | Record<string, Array<CompletionAnnotation> | Map<string, string | number>>
     | undefined;
 };
 
+function listAnnotationDisplay(listAnnotation: Array<CompletionAnnotation>) {
+  return (
+    <div>
+      {listAnnotation.map((annotation, idx) => (
+        <div key={idx}>
+          {annotation.error && (
+            <div>
+              <h3 className="ml-1">Error</h3>
+              <Preview value={annotation["error"]} />{" "}
+            </div>
+          )}
+          {annotation.text && (
+            <div>
+              <h3 className="ml-1">Text</h3>
+              <Preview value={annotation["text"]} />{" "}
+            </div>
+          )}
+          {annotation.media_object && (
+            <MediaObjectDisplay mediaObject={annotation["media_object"]} />
+          )}
+        </div>
+      ))}
+    </div>
+  );
+}
+
+function dictAnnotationDisplay(
+  dictAnnotation: Record<string, string | number>,
+) {
+  return (
+    <div>
+      {Object.entries(dictAnnotation).map(([key, value]) => (
+        <div>
+          <h3 className="ml-1">{key}</h3>
+          <Preview value={value.toString()} />
+        </div>
+      ))}
+    </div>
+  );
+}
+
 export default function AnnotationDisplay({ predictionAnnotations }: Props) {
   return (
     <div>
@@ -17,27 +61,9 @@ export default function AnnotationDisplay({ predictionAnnotations }: Props) {
               <h3>
                 <strong>{key}</strong>
               </h3>
-              {value.map((annotation, idx) => (
-                <div key={idx}>
-                  {annotation.error && (
-                    <div>
-                      <h3 className="ml-1">Error</h3>
-                      <Preview value={annotation["error"]} />{" "}
-                    </div>
-                  )}
-                  {annotation.text && (
-                    <div>
-                      <h3 className="ml-1">Text</h3>
-                      <Preview value={annotation["text"]} />{" "}
-                    </div>
-                  )}
-                  {annotation.media_object && (
-                    <MediaObjectDisplay
-                      mediaObject={annotation["media_object"]}
-                    />
-                  )}
-                </div>
-              ))}
+              {Array.isArray(value)
+                ? listAnnotationDisplay(value)
+                : dictAnnotationDisplay(value)}
             </div>
           ))
         : null}