Upload cosine_magnitude_audit_v2.py

cosine_magnitude_audit_v2.py

@@ -0,0 +1,104 @@
+import os
+import torch
+from safetensors.torch import safe_open
+import yaml
+
+# --- CONFIGURATION ---
+YAML_PATH = "B:/24B/C82/mergekit_config.yml"
+FINAL_MERGE_DIR = "B:/24B/C82"
+LAYERS_TO_SCAN = [
+    "model.layers.30.mlp.down_proj.weight"
+]
+# ---------------------
+
+def load_tensor(model_dir, tensor_name):
+    """Finds and loads a tensor from a directory of safetensors."""
+    for file in os.listdir(model_dir):
+        if file.endswith(".safetensors"):
+            filepath = os.path.join(model_dir, file)
+            with safe_open(filepath, framework="pt", device="cpu") as f:
+                if tensor_name in f.keys():
+                    return f.get_tensor(tensor_name).float()
+    raise ValueError(f"Tensor {tensor_name} not found in {model_dir}")
+
+def get_latent_metrics(tv):
+    """Calculates structural metrics of the task vector."""
+    # 1. Global Energy (Frobenius Norm)
+    f_norm = torch.linalg.norm(tv).item()
+    
+    # 2. Spectral Norm (Intensity of dominant feature)
+    # Measures the largest singular value (L2 operator norm)
+    s_norm = torch.linalg.norm(tv, ord=2).item()
+    
+    # 3. Stable Rank (Knowledge Breadth / Complexity)
+    # Calculation: ||A||_F^2 / ||A||_2^2
+    # Higher = more diverse features preserved; Lower = collapsed/overfit
+    s_rank = (f_norm**2) / (s_norm**2) if s_norm > 0 else 0
+    
+    return f_norm, s_norm, s_rank
+
+def main():
+    print("Loading YAML config...")
+    with open(YAML_PATH, 'r') as f:
+        config = yaml.safe_load(f)
+    
+    base_path = config['base_model']
+    donor_paths = [m['model'] for m in config['models']]
+    
+    print(f"\nScanning {len(LAYERS_TO_SCAN)} MLP layers for structural influence...\n")
+    
+    for layer in LAYERS_TO_SCAN:
+        print(f"--- Layer: {layer} ---")
+        try:
+            # Load Base and Final
+            base_w = load_tensor(base_path, layer)
+            final_w = load_tensor(FINAL_MERGE_DIR, layer)
+            final_tv = final_w - base_w
+            
+            # Calculate metrics for the Final Merged Model
+            final_f, final_s, final_rk = get_latent_metrics(final_tv)
+            
+            print(f"Final Model Latent Metrics:")
+            print(f"  Energy (Frobenius):   {final_f:.4f}")
+            print(f"  Intensity (Spectral): {final_s:.4f}")
+            print(f"  Complexity (S-Rank):  {final_rk:.2f}")
+            print("")
+
+            results = []
+            for donor in donor_paths:
+                donor_w = load_tensor(donor, layer)
+                donor_tv = donor_w - base_w
+                
+                # Get structural metrics for donor
+                d_f, d_s, d_rk = get_latent_metrics(donor_tv)
+                
+                # Calculate Cosine Similarity (Alignment)
+                cos_sim = torch.nn.functional.cosine_similarity(
+                    final_tv.flatten(), donor_tv.flatten(), dim=0
+                ).item()
+                
+                name = donor.split("/")[-1][:40]
+                results.append({
+                    'name': name,
+                    'sim': cos_sim,
+                    'energy': d_f,
+                    'intens': d_s,
+                    'cmplx': d_rk
+                })
+            
+            # Sort by highest alignment
+            results.sort(key=lambda x: x['sim'], reverse=True)
+            
+            print(f"{'Donor Model':<42} | {'Align':<8} | {'Energy':<8} | {'Intens':<8} | {'Cmplx'}")
+            print("-" * 90)
+            for r in results:
+                # Energy is relative to the final model (e.g. 0.90x)
+                rel_energy = r['energy'] / final_f if final_f > 0 else 0
+                print(f"{r['name']:<42} | {r['sim']:>7.3f} | {rel_energy:>7.2f}x | {r['intens']:>8.3f} | {r['cmplx']:>6.2f}")
+            print("\n")
+            
+        except Exception as e:
+            print(f"Skipping layer due to error: {e}")
+
+if __name__ == "__main__":
+    main()