Update app.py

app.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 import gradio as gr
 import numpy as np
@@ -113,11 +113,14 @@ def highlight_token(token: str, score: float):
 
 
 def create_highlighted_text(
-    label: str, tokens2scores: List[Tuple[str, float]], mean_surprisal: float
+    label: str,
+    tokens2scores: List[Tuple[str, float]],
+    mean_surprisal: Optional[float] = None,
 ):
-    highlighted_text: str = (
-        "<h2><b>" + label + f"</b>（サプライザル平均値: {mean_surprisal}）</h2>"
-    )
+    if mean_surprisal is None:
+        highlighted_text = "<h2><b>" + label + "</b></h2>"
+    else:
+        highlighted_text = "<h2><b>" + label + f"</b>（サプライザル平均値: {mean_surprisal}）</h2>"
     for token, score in tokens2scores:
         highlighted_text += highlight_token(token, score)
     return highlighted_text
@@ -184,7 +187,9 @@ def main(input_text: str) -> Tuple[str, str, str]:
     diff_tokens2surprisal = calculate_surprisal_diff(
         tokens2surprisal, baseline_tokens2surprisal, 100.0
     )
-    diff_highlighted_text = create_highlighted_text("学習前後の差分", diff_tokens2surprisal)
+    diff_highlighted_text = create_highlighted_text(
+        "学習前後の差分", diff_tokens2surprisal, None
+    )
     return (
         baseline_highlighted_text,
         highlighted_text,