231 lines
9.3 KiB
Python
231 lines
9.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Analyzes positional embeddings quality.
|
|
Checks if positions 77-247 are distinguishable enough
|
|
for the model to understand word order.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from safetensors import safe_open
|
|
import sys
|
|
import os
|
|
|
|
|
|
def analyze(checkpoint_path):
|
|
print(f"\nAnalyzing: {os.path.basename(checkpoint_path)}")
|
|
print("=" * 70)
|
|
|
|
# Load positional embeddings
|
|
pos_embs = {}
|
|
with safe_open(checkpoint_path, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
if "position" in key.lower() and "embedding" in key.lower():
|
|
tensor = f.get_tensor(key)
|
|
if tensor.dim() == 2 and tensor.shape[0] > 77:
|
|
pos_embs[key] = tensor.float()
|
|
|
|
if not pos_embs:
|
|
print("No extended positional embeddings found!")
|
|
return
|
|
|
|
for key, emb in pos_embs.items():
|
|
n_pos, dim = emb.shape
|
|
print(f"\n{'='*70}")
|
|
print(f"Key: {key}")
|
|
print(f"Shape: ({n_pos}, {dim})")
|
|
print(f"{'='*70}")
|
|
|
|
# ============================================
|
|
# Test 1: Pairwise cosine similarity
|
|
# ============================================
|
|
print(f"\n--- Test 1: Cosine Similarity Between Positions ---")
|
|
|
|
# Normalize
|
|
emb_norm = emb / (torch.norm(emb, dim=1, keepdim=True) + 1e-8)
|
|
sim_matrix = emb_norm @ emb_norm.T # (n_pos, n_pos)
|
|
|
|
# Original zone (0-76)
|
|
orig_sim = sim_matrix[:77, :77]
|
|
orig_offdiag = orig_sim[~torch.eye(77, dtype=torch.bool)]
|
|
print(f"\n Original positions (0-76):")
|
|
print(f" Mean self-similarity: {orig_offdiag.mean():.4f}")
|
|
print(f" Min self-similarity: {orig_offdiag.min():.4f}")
|
|
print(f" Max self-similarity: {orig_offdiag.max():.4f}")
|
|
print(f" Std: {orig_offdiag.std():.4f}")
|
|
|
|
# New zone (77-247)
|
|
if n_pos > 77:
|
|
new_sim = sim_matrix[77:, 77:]
|
|
new_offdiag = new_sim[~torch.eye(n_pos - 77, dtype=torch.bool)]
|
|
print(f"\n New positions (77-{n_pos-1}):")
|
|
print(f" Mean self-similarity: {new_offdiag.mean():.4f}")
|
|
print(f" Min self-similarity: {new_offdiag.min():.4f}")
|
|
print(f" Max self-similarity: {new_offdiag.max():.4f}")
|
|
print(f" Std: {new_offdiag.std():.4f}")
|
|
|
|
if new_offdiag.mean() > 0.95:
|
|
print(f"\n *** WARNING: New positions are too similar! ***")
|
|
print(f" *** Model may not distinguish word order ***")
|
|
elif new_offdiag.mean() > 0.85:
|
|
print(f"\n ** CAUTION: Similarity is high **")
|
|
print(f" ** Word order may be weakly encoded **")
|
|
else:
|
|
print(f"\n OK: Positions are distinguishable")
|
|
|
|
# Cross zone (original vs new)
|
|
if n_pos > 77:
|
|
cross_sim = sim_matrix[:77, 77:]
|
|
print(f"\n Cross-zone (original vs new):")
|
|
print(f" Mean similarity: {cross_sim.mean():.4f}")
|
|
|
|
# ============================================
|
|
# Test 2: Per-position norm
|
|
# ============================================
|
|
print(f"\n--- Test 2: Position Norms ---")
|
|
|
|
norms = torch.norm(emb, dim=1)
|
|
print(f"\n Original (0-76):")
|
|
print(f" Mean norm: {norms[:77].mean():.4f}")
|
|
print(f" Std norm: {norms[:77].std():.4f}")
|
|
|
|
if n_pos > 77:
|
|
print(f"\n New (77-{n_pos-1}):")
|
|
print(f" Mean norm: {norms[77:].mean():.4f}")
|
|
print(f" Std norm: {norms[77:].std():.4f}")
|
|
print(f" Norm ratio (new/orig): {norms[77:].mean() / norms[:77].mean():.4f}")
|
|
|
|
if norms[77:].std() < norms[:77].std() * 0.1:
|
|
print(f"\n *** WARNING: New positions have collapsed norms ***")
|
|
print(f" *** They all look the same to the model ***")
|
|
|
|
# ============================================
|
|
# Test 3: Consecutive distance
|
|
# ============================================
|
|
print(f"\n--- Test 3: Distance Between Consecutive Positions ---")
|
|
|
|
diffs = torch.norm(emb[1:] - emb[:-1], dim=1)
|
|
|
|
print(f"\n Original (0-76):")
|
|
print(f" Mean step: {diffs[:76].mean():.4f}")
|
|
print(f" Std step: {diffs[:76].std():.4f}")
|
|
|
|
if n_pos > 78:
|
|
print(f"\n New (77-{n_pos-1}):")
|
|
print(f" Mean step: {diffs[77:].mean():.4f}")
|
|
print(f" Std step: {diffs[77:].std():.4f}")
|
|
|
|
boundary = diffs[76].item()
|
|
orig_mean = diffs[:76].mean().item()
|
|
print(f"\n Boundary (76->77): {boundary:.4f}")
|
|
print(f" Original mean step: {orig_mean:.4f}")
|
|
print(f" Boundary ratio: {boundary / orig_mean:.4f}")
|
|
|
|
if diffs[77:].mean() < diffs[:76].mean() * 0.1:
|
|
print(f"\n *** WARNING: Steps between new positions are tiny ***")
|
|
print(f" *** Positions are collapsing together ***")
|
|
|
|
# ============================================
|
|
# Test 4: Uniqueness (rank)
|
|
# ============================================
|
|
print(f"\n--- Test 4: Effective Dimensionality ---")
|
|
|
|
_, S, _ = torch.linalg.svd(emb)
|
|
total = (S ** 2).sum()
|
|
cumulative = torch.cumsum(S ** 2, dim=0) / total
|
|
|
|
rank_95 = int((cumulative < 0.95).sum().item()) + 1
|
|
rank_99 = int((cumulative < 0.99).sum().item()) + 1
|
|
|
|
print(f" Full embedding ({n_pos} positions):")
|
|
print(f" Rank for 95% energy: {rank_95}")
|
|
print(f" Rank for 99% energy: {rank_99}")
|
|
|
|
if n_pos > 77:
|
|
_, S_orig, _ = torch.linalg.svd(emb[:77])
|
|
total_o = (S_orig ** 2).sum()
|
|
cum_o = torch.cumsum(S_orig ** 2, dim=0) / total_o
|
|
rank_95_o = int((cum_o < 0.95).sum().item()) + 1
|
|
|
|
_, S_new, _ = torch.linalg.svd(emb[77:])
|
|
total_n = (S_new ** 2).sum()
|
|
cum_n = torch.cumsum(S_new ** 2, dim=0) / total_n
|
|
rank_95_n = int((cum_n < 0.95).sum().item()) + 1
|
|
|
|
print(f" Original (0-76): rank_95 = {rank_95_o}")
|
|
print(f" New (77-{n_pos-1}): rank_95 = {rank_95_n}")
|
|
|
|
if rank_95_n < rank_95_o // 2:
|
|
print(f"\n *** WARNING: New positions have low rank ***")
|
|
print(f" *** They lack diversity compared to original ***")
|
|
|
|
# ============================================
|
|
# Test 5: Position discrimination ability
|
|
# ============================================
|
|
print(f"\n--- Test 5: Can nearby positions be distinguished? ---")
|
|
|
|
# For each window of 10 positions, check if they're unique
|
|
window = 10
|
|
for start in [0, 37, 70, 77, 120, 170, 220]:
|
|
end = min(start + window, n_pos)
|
|
if start >= n_pos:
|
|
break
|
|
chunk = emb[start:end]
|
|
chunk_norm = chunk / (torch.norm(chunk, dim=1, keepdim=True) + 1e-8)
|
|
chunk_sim = chunk_norm @ chunk_norm.T
|
|
mask = ~torch.eye(end - start, dtype=torch.bool)
|
|
avg_sim = chunk_sim[mask].mean().item()
|
|
min_sim = chunk_sim[mask].min().item()
|
|
|
|
zone = "ORIG" if start < 77 else "NEW "
|
|
status = "OK" if avg_sim < 0.95 else "WARN" if avg_sim < 0.99 else "BAD"
|
|
print(f" [{zone}] Positions {start:3d}-{end-1:3d}: "
|
|
f"avg_sim={avg_sim:.4f} min_sim={min_sim:.4f} [{status}]")
|
|
|
|
# ============================================
|
|
# Summary
|
|
# ============================================
|
|
print(f"\n{'='*70}")
|
|
print(f"SUMMARY")
|
|
print(f"{'='*70}")
|
|
|
|
issues = []
|
|
|
|
if n_pos > 77:
|
|
new_offdiag = sim_matrix[77:, 77:]
|
|
mask = ~torch.eye(n_pos - 77, dtype=torch.bool)
|
|
mean_sim = new_offdiag[mask].mean().item()
|
|
|
|
if mean_sim > 0.95:
|
|
issues.append("Positions 77+ are nearly identical — word order lost")
|
|
elif mean_sim > 0.85:
|
|
issues.append("Positions 77+ are quite similar — word order weakly encoded")
|
|
|
|
if norms[77:].std() < norms[:77].std() * 0.1:
|
|
issues.append("Position norms have collapsed — low diversity")
|
|
|
|
if diffs[77:].mean() < diffs[:76].mean() * 0.1:
|
|
issues.append("Steps between positions are tiny — positions collapsed")
|
|
|
|
if issues:
|
|
print(f"\n ISSUES FOUND:")
|
|
for issue in issues:
|
|
print(f" ! {issue}")
|
|
print(f"\n Recommendation: try different extension method")
|
|
else:
|
|
print(f"\n ALL CHECKS PASSED")
|
|
print(f" Positions are distinguishable — word order should be understood")
|
|
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python analyze_positions.py <checkpoint.safetensors>")
|
|
print("\nAnalyze original and converted:")
|
|
print(" python analyze_positions.py model.safetensors")
|
|
print(" python analyze_positions.py model_long248.safetensors")
|
|
sys.exit(1)
|
|
|
|
analyze(sys.argv[1])
|