Add overlap-aware turn prediction objectives

streaming_turn_model_directly.py

@@ -1540,7 +1540,7 @@ def main(
         "verification": result.verification,
     }
     if args.head_checkpoint is not None:
-        head = core._load_head(
+        head, head_loss_name = core._load_head(
             args.head_checkpoint, core._hidden_size(model), model.device
         )
         with torch.inference_mode():
@@ -1548,7 +1548,11 @@ def main(
                 result.states.to(model.device, dtype=model.dtype)
             ).float().cpu()
         payload["logits"] = logits
-        payload["probabilities"] = logits.softmax(dim=-1)
+        payload["probabilities"] = core.marginal_probabilities_from_logits(
+            logits,
+            head_loss_name,
+        )
+        payload["head_loss"] = head_loss_name
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
     torch.save(payload, args.output)