Add overlap-aware turn prediction objectives

streaming_turn_model.py

@@ -44,7 +44,9 @@ import torch
 from easycom_head import (
     TurnPredictionHead,
     checkpoint_head_config,
+    checkpoint_loss_name,
     egospeak_multihorizon_loss,
+    marginal_probabilities_from_logits,
     miniroad_multihorizon_loss,
     multihorizon_cross_entropy,
 )
@@ -2141,10 +2143,11 @@ def _load_head(
     checkpoint: Path,
     hidden_size: int,
     device: torch.device,
-) -> TurnPredictionHead:
+) -> tuple[TurnPredictionHead, str]:
     payload = torch.load(checkpoint, map_location="cpu", weights_only=True)
     if isinstance(payload, dict) and "head" in payload:
         config = checkpoint_head_config(payload)
+        loss_name = checkpoint_loss_name(payload)
         if int(config["hidden_size"]) != hidden_size:
             raise ValueError(
                 f"Head hidden size {config['hidden_size']} != model {hidden_size}"
@@ -2154,8 +2157,9 @@ def _load_head(
     else:
         head = TurnPredictionHead(hidden_size=hidden_size)
         state_dict = payload
+        loss_name = "lstr-mce"
     head.load_state_dict(state_dict)
-    return head.to(device).eval()
+    return head.to(device).eval(), loss_name
 
 
 def parse_args() -> argparse.Namespace:
@@ -2393,11 +2397,19 @@ def main() -> None:
         "verification": result.verification,
     }
     if args.head_checkpoint is not None:
-        head = _load_head(args.head_checkpoint, _hidden_size(model), model.device)
+        head, head_loss_name = _load_head(
+            args.head_checkpoint,
+            _hidden_size(model),
+            model.device,
+        )
         with torch.inference_mode():
             logits = head(result.states.to(model.device, dtype=model.dtype)).float().cpu()
         payload["logits"] = logits
-        payload["probabilities"] = logits.softmax(dim=-1)
+        payload["probabilities"] = marginal_probabilities_from_logits(
+            logits,
+            head_loss_name,
+        )
+        payload["head_loss"] = head_loss_name
 
     if args.num_steps is not None and len(result.times) != args.num_steps:
         raise RuntimeError(