Update app.py

app.py

@@ -164,7 +164,7 @@ class ModelRunner:
         tokens = [self.tokenizer.decode(i) for i in input_ids]
         return pd.DataFrame({"token": tokens, "pred": eou_probs_masked})
 
-    def make_styled_df(self, df: pd.DataFrame, thresh: float, cmap="coolwarm") -> str:
+    def make_styled_df(self, df: pd.DataFrame, cmap="coolwarm") -> str:
         EPS = 1e-12
         thresh = self.thresh
         
@@ -201,7 +201,7 @@ class ModelRunner:
         )
         return fmt.to_html()
 
-    def generate_highlighted_text(self, text: str, threshold: float):
+    def generate_highlighted_text(self, text: str):
         """Returns: (highlighted_list, styled_html) for Gradio"""
         eps = 1e-12
         threshold = self.thresh
@@ -222,7 +222,7 @@ class ModelRunner:
         if pd.notna(max_abs_score) and max_abs_score > 0:
             df.score = df.score / (max_abs_score * 1.5)
 
-        styled_df = self.make_styled_df(df[["token", "pred"]], thresh=threshold)
+        styled_df = self.make_styled_df(df[["token", "pred"]])
         return list(zip(df.token, df.score)), styled_df