Allow dicts to be displayed in AnnotationsDisplay without a enclosing…

… list (stanford-crfm#2700)
xuwangyinx · Jun 23, 2024 · e6a69d6 · e6a69d6
1 parent 808ae62
commit e6a69d6
Showing 1 changed file with 48 additions and 22 deletions.
diff --git 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}