Update app.py

app.py

@@ -66,7 +66,7 @@ def infer(model, replay_file,
     predictions = []
     it = trange(len(serialized_states), desc="Running model")
     for i in range(0, len(serialized_states), bs):
-        batch = (serialized_states[i:i + bs].clone().to(DEVICE), 
+        batch = (serialized_states[i:i + bs].clone().to(DEVICE),
                  serialized_scoreboards[i:i + bs].clone().to(DEVICE))
         if nullify_goal_difference or ignore_ties:
             batch[1][:, SB_BLUE_SCORE] = 0
@@ -115,13 +115,17 @@ def infer(model, replay_file,
     preds = preds[main_cols + [c for c in preds.columns if c not in main_cols]]
     # Set index name
     preds.index.name = "Frame"
-    if ignore_ties:
-        tie_probs = preds["Tie"]
+    remove_ties_mask = is_ot if not ignore_ties else torch.ones(len(preds), dtype=torch.bool)
+    if remove_ties_mask.any():
+        tie_probs = preds[remove_ties_mask, "Tie"]
         q = (1 - tie_probs)
-        preds = preds.drop("Tie", axis=1)
         for c in preds.columns:
             if c.startswith("Blue") or c.startswith("Orange"):
-                preds[c] /= q
+                preds[remove_ties_mask, c] /= q
+        if ignore_ties:
+            preds = preds.drop("Tie", axis=1)
+        else:
+            preds[remove_ties_mask, "Tie"] = 0.0
 
     return preds
 
@@ -242,21 +246,36 @@ def plot_plotly(preds: pd.DataFrame):
     return fig
 
 
+DESCRIPTION = """
+# Next Goal Predictor
+Upload a replay file to get a plot of the next goal prediction.
+
+The model is trained on about 50,000 hours of SSL and RLCS replays in 1s, 2s, and 3s.
+
+It predicts the probability of each team scoring the next goal, in addition to ties.
+It also predicts the probability that each team will score in 1s, 2s, etc. up to 60+ seconds.
+
+The plot only shows the team predictions, but you can download the full predictions. 
+""".strip()
+
+RADIO_OPTIONS = [
+    "**Nullify goal difference**<br>Makes the model think the goal difference is always 0, so it doesn't have a bias towards one team.",
+    "**Ignore ties**<br>Makes the model pretend every situation is an overtime (e.g. ties are impossible)."
+]
+
+
 def gradio_app():
     import gradio as gr
 
     with TemporaryDirectory() as temp_dir:
         with gr.Blocks() as demo:
-            gr.Markdown("# Next Goal Predictor")
-            gr.Markdown("Upload a replay file to get a plot of the next goal prediction.<br>"
-                        "The model is trained on about 50 000 hours of SSL and RLCS replays in 1s, 2s, and 3s.")
+            gr.Markdown(DESCRIPTION)
 
             # Use gr.Column to stack components vertically
             with gr.Column():
                 file_input = gr.File(label="Upload Replay File", type="filepath", file_types=[".replay"])
                 checkboxes = gr.Radio(label="Options", choices=["Nullify goal difference", "Ignore ties"], type="index",
-                                     info="You can choose to make the model think the goal difference is always 0 (so it doesn't prefer one team), "
-                                          "or make it pretend every situation is an overtime (e.g. ties are impossible).")
+                                      info=RADIO_INFO)
                 submit_button = gr.Button("Generate Predictions")
                 plot_output = gr.Plot(label="Predictions")
                 download_button = gr.DownloadButton("Download Predictions", visible=False)
@@ -266,7 +285,7 @@ def gradio_app():
                 nullify_goal_difference = radio_option == 0
                 ignore_ties = radio_option == 1
                 print(f"Processing file: {replay_file}")
-                
+
                 replay_stem = os.path.splitext(os.path.basename(replay_file))[0]
                 postfix = ""
                 if nullify_goal_difference:
@@ -278,9 +297,9 @@ def gradio_app():
                     print(f"Predictions file already exists: {preds_file}")
                     preds = pd.read_csv(preds_file, dtype={"Touch": str})
                 else:
-                    preds = infer(MODEL, replay_file, 
-                                 nullify_goal_difference=nullify_goal_difference,
-                                 ignore_ties=ignore_ties)
+                    preds = infer(MODEL, replay_file,
+                                  nullify_goal_difference=nullify_goal_difference,
+                                  ignore_ties=ignore_ties)
                 plt = plot_plotly(preds)
                 print(f"Plot generated for file: {replay_file}")
                 preds.to_csv(preds_file)