Simplify result table labels

app.py

@@ -36,6 +36,14 @@ os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 ROOT_DIR = Path(__file__).resolve().parent
 CONFIG_PATH = ROOT_DIR / "configs" / "medical_vqa.yaml"
 VARIANT_ORDER = ["A1", "A2", "B1", "B2", "DPO", "PPO"]
+MODEL_DISPLAY_NAMES = {
+    "A1": "A1 LSTM",
+    "A2": "A2 Transformer",
+    "B1": "B1 Zero-shot",
+    "B2": "B2 Fine-tuned",
+    "DPO": "DPO Alignment",
+    "PPO": "PPO RL refinement",
+}
 HF_MODEL_REPOS = {
     "A1": "SpringWang08/medical-vqa-a1",
     "A2": "SpringWang08/medical-vqa-a2",
@@ -323,7 +331,9 @@ async def _predict_variant(variant: str, question: str, image: Image.Image) -> d
         )
         return {
             "model": variant,
+            "Model": MODEL_DISPLAY_NAMES.get(variant, variant),
             "prediction": rewritten,
+            "Prediction": rewritten,
             "prediction_before_rewrite": out["prediction"],
             "raw": out["prediction_raw"],
             "answer_used_for_rewrite": answer_for_rewrite,
@@ -334,7 +344,9 @@ async def _predict_variant(variant: str, question: str, image: Image.Image) -> d
     except Exception as exc:
         return {
             "model": variant,
+            "Model": MODEL_DISPLAY_NAMES.get(variant, variant),
             "prediction": "",
+            "Prediction": "",
             "prediction_before_rewrite": "",
             "raw": "",
             "answer_used_for_rewrite": "",
@@ -362,7 +374,7 @@ def predict_all(image: Image.Image, question: str, selected_models: list[str]) -
         return rows
 
     rows = asyncio.run(_run())
-    return pd.DataFrame(rows)
+    return pd.DataFrame(rows)[["Model", "Prediction"]]
 
 
 CSS = """
@@ -390,14 +402,8 @@ with gr.Blocks(css=CSS, title="Medical VQA Compare") as demo:
             output_table = gr.Dataframe(
                 label="Kết quả",
                 headers=[
-                    "model",
-                    "prediction",
-                    "prediction_before_rewrite",
-                    "raw",
-                    "answer_used_for_rewrite",
-                    "checkpoint",
-                    "latency_ms",
-                    "status",
+                    "Model",
+                    "Prediction",
                 ],
                 wrap=True,
             )