Multi-head attention

# 10. dersteki TEK başlıklı self-attention, "o -> Kedi" ilişkisini GÜZEL yakaladı
# (ağırlık 0.517). Ama AYNI başlık, "çıktı" kelimesinin NEREYE (masaya) işaret
# ettiğini yakalayamadı -- ağırlıkları neredeyse DÜZ (uniform):
TOKENLER = ["Kedi", "masaya", "çıktı", "çünkü", "o", "yorulmuştu"]
tek_baslik_ciktI_agirliklari = [0.136, 0.170, 0.204, 0.177, 0.139, 0.174]
print("10. dersteki TEK başlığın 'çıktı' satırı (10. dersten alınan gerçek sonuç):")
for tok, a in zip(TOKENLER, tek_baslik_ciktI_agirliklari):
    print(f"  {tok:<12} ağırlık={a:.3f}")
print(f"\nEn yüksek ({max(tek_baslik_ciktI_agirliklari):.3f}) ile en düşük ({min(tek_baslik_ciktI_agirliklari):.3f}) arasındaki fark ÇOK küçük --")
print("bu başlık 'eylem-konum' (çıktı -> masaya) ilişkisini YAKALAYAMADI, çünkü boyutları")
print("SADECE 'coreference' (o -> Kedi) ilişkisine göre AYARLANMIŞTI. TEK başlık, TEK ilişki türü demek.")

# Çözüm: AYNI anda BİRDEN FAZLA başlık çalıştır, HER biri FARKLI bir alt-uzayda
# (subspace) FARKLI bir ilişki türü arasın. D=8 boyutu, H=2 başlığa (her biri 4
# boyutlu) BÖLÜNÜYOR.
H, D_TOPLAM, D_BAS = 2, 8, 4

# Head 1 (boyut 0:4) -- coreference ilişkisi için 10. dersteki AYNI ağırlıklar (seed=2426).
torch.manual_seed(2426)
Wq1, Wk1, Wv1 = torch.randn(D_BAS, D_BAS) * 0.5, torch.randn(D_BAS, D_BAS) * 0.5, torch.randn(D_BAS, D_BAS) * 0.5
# Head 2 (boyut 4:8) -- eylem-konum ilişkisi için AYRI, taranmış bir seed (seed=9678).
torch.manual_seed(9678)
Wq2, Wk2, Wv2 = torch.randn(D_BAS, D_BAS) * 0.5, torch.randn(D_BAS, D_BAS) * 0.5, torch.randn(D_BAS, D_BAS) * 0.5

# Gerçek Transformer'larda TEK bir Wq (D×D) kullanılır; biz burada ÖĞRETİCİ amaçlı
# BLOK-KÖŞEGEN (block-diagonal) bir Wq inşa ediyoruz -- bu, "her başlık kendi
# alt-uzayını görsün" fikrini AÇIKÇA gösteriyor (gerçek modellerde bu uzmanlaşma
# EĞİTİMLE, dolaylı olarak ortaya çıkar).
def blok_kosegen(A, B):
    ust = torch.cat([A, torch.zeros(D_BAS, D_BAS)], dim=1)
    alt = torch.cat([torch.zeros(D_BAS, D_BAS), B], dim=1)
    return torch.cat([ust, alt], dim=0)

Wq_tam = blok_kosegen(Wq1, Wq2)
Wk_tam = blok_kosegen(Wk1, Wk2)
Wv_tam = blok_kosegen(Wv1, Wv2)
print(f"Wq_tam şekli: {tuple(Wq_tam.shape)} -- {H} başlık x {D_BAS} boyut = {D_TOPLAM} toplam boyut.")

# İki ilişkiyi de barındıran 8 boyutlu gömme: ilk 4 boyut 10. dersteki coreference
# eksenleri, son 4 boyut YENİ bir eylem-konum ekseni.
X = torch.tensor([
    [1.5, 0.0, 0.0, 0.0,   0.0, 0.0, 0.0, 0.0],   # Kedi
    [0.0, 1.5, 0.0, 0.0,   1.4, 0.0, 0.0, 0.0],   # masaya
    [0.3, 0.3, 1.2, 0.0,   1.3, 0.2, 0.0, 0.0],   # çıktı
    [0.0, 0.0, 0.0, 1.0,   0.0, 0.0, 0.0, 0.8],   # çünkü
    [1.3, 0.0, 0.0, 0.1,   0.0, 0.0, 0.0, 0.0],   # o
    [0.2, 0.0, 0.9, 0.1,   0.0, 0.0, 1.0, 0.1],   # yorulmuştu
])
Q = X @ Wq_tam
K = X @ Wk_tam
V = X @ Wv_tam
print(f"\nX şekli: {tuple(X.shape)} -> Q, K, V şekli (üçü de): {tuple(Q.shape)}")

# Q, K, V'yi (N, H, D_BAS) şeklinde YENİDEN ŞEKİLLENDİR, HER başlığı AYRI AYRI çalıştır.
N = X.shape[0]
Q_bas = Q.view(N, H, D_BAS).transpose(0, 1)  # (H, N, D_BAS)
K_bas = K.view(N, H, D_BAS).transpose(0, 1)
V_bas = V.view(N, H, D_BAS).transpose(0, 1)

baslik_ciktilari = []
baslik_agirliklari = []
for h in range(H):
    cikti_h, agirlik_h = dikkat(Q_bas[h], K_bas[h], V_bas[h])
    baslik_ciktilari.append(cikti_h)
    baslik_agirliklari.append(agirlik_h)

ciktI_idx, o_idx = TOKENLER.index("çıktı"), TOKENLER.index("o")
print("\nBaşlık 1 (coreference alt-uzayı) -- 'o' satırı:")
for tok, a in zip(TOKENLER, baslik_agirliklari[0][o_idx].tolist()):
    print(f"  {tok:<12} ağırlık={a:.3f}")
print("\nBaşlık 2 (eylem-konum alt-uzayı) -- 'çıktı' satırı:")
for tok, a in zip(TOKENLER, baslik_agirliklari[1][ciktI_idx].tolist()):
    print(f"  {tok:<12} ağırlık={a:.3f}")

en_yuksek_1 = TOKENLER[baslik_agirliklari[0][o_idx].argmax().item()]
en_yuksek_2 = TOKENLER[baslik_agirliklari[1][ciktI_idx].argmax().item()]
print(f"\nBaşlık 1: 'o' en çok '{en_yuksek_1}'e bakıyor. Başlık 2: 'çıktı' en çok '{en_yuksek_2}'ya bakıyor.")
print("AYNI cümle, AYNI anda, İKİ FARKLI ilişki türü -- her başlık kendi 'uzmanlığında'.")

# Başlıkları BİRLEŞTİR (concat) -- bu, çok-başlıklı attention'ın çıktısı.
coklu_baslik_cikti = torch.cat(baslik_ciktilari, dim=-1)
print(f"\nBirleştirilmiş çıktı şekli: {tuple(coklu_baslik_cikti.shape)} (= {H} x {D_BAS})")

# Blok-köşegen TAM matrislerle (X @ Wq_tam, ardından reshape) hesaplamak ile,
# HER başlığı kendi 4 boyutluk diliminde AYRI AYRI hesaplamak AYNI sonucu vermeli mi?
Q1_ayri, K1_ayri, V1_ayri = X[:, :4] @ Wq1, X[:, :4] @ Wk1, X[:, :4] @ Wv1
Q2_ayri, K2_ayri, V2_ayri = X[:, 4:] @ Wq2, X[:, 4:] @ Wk2, X[:, 4:] @ Wv2
cikti1_ayri, _ = dikkat(Q1_ayri, K1_ayri, V1_ayri)
cikti2_ayri, _ = dikkat(Q2_ayri, K2_ayri, V2_ayri)
cikti_ayri_birlesik = torch.cat([cikti1_ayri, cikti2_ayri], dim=-1)

fark = (coklu_baslik_cikti - cikti_ayri_birlesik).abs().max().item()
print(f"\nBlok-köşegen TAM matrisle vs AYRI AYRI hesaplanan başlıklar arasındaki fark: {fark:.2e}")
print(f"torch.allclose: {torch.allclose(coklu_baslik_cikti, cikti_ayri_birlesik, atol=1e-6)}")
print("Bu, 'blok-köşegen Wq = her başlığı kendi alt-uzayında çalıştırmakla AYNI' iddiasını DOĞRULUYOR.")