Add learnable last-four-layer query fusion

infer_turn_head.py

@@ -283,6 +283,11 @@ def main() -> None:
                 f"{name} feature dim is {data.hidden_size}, checkpoint expects "
                 f"{head_config['hidden_size']}"
             )
+        if data.num_input_layers != int(head_config["num_input_layers"]):
+            raise ValueError(
+                f"{name} has {data.num_input_layers} feature layers, checkpoint "
+                f"expects {head_config['num_input_layers']}"
+            )
         probabilities = _predict_states(
             head,
             data.states,
@@ -323,14 +328,29 @@ def main() -> None:
             "checkpoint_epoch": checkpoint_payload.get("epoch"),
             "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
             "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
+            "layer_mixture_weights_oldest_to_newest": (
+                head.layer_mixture_weights().detach().float().cpu().tolist()
+                if head.layer_mixture_weights() is not None
+                else None
+            ),
         }
     )
     metrics_path = args.output_dir / "metrics.json"
     metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
     _print_metrics(metrics)
+    if metrics["layer_mixture_weights_oldest_to_newest"] is not None:
+        print(
+            "\nLayer mixture weights (oldest->newest): "
+            + ", ".join(
+                f"{value:.6f}"
+                for value in metrics[
+                    "layer_mixture_weights_oldest_to_newest"
+                ]
+            )
+        )
     print(f"\nmetrics: {metrics_path}")
     print(f"EgoSpeak-compatible predictions: {prediction_dir}")
 
 
 if __name__ == "__main__":
-    main()
+    main()