{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n",
      "Iteration 0, Current Loss: 132.482483\n",
      "Iteration 100, Current Loss: 0.000000\n",
      "Loss is small enough. Breaking early.\n",
      "find 512 vectors in dim 128 with max_abs_dot = 0.12004002183675766\n",
      "torch.Size([512, 128])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training 1 Layer 128 Dim Delta Kernel round+round: 100%|██████████| 2000/2000 [04:19<00:00,  7.69it/s, loss=0.2930, acc=0.9157]\n",
      "Training 1 Layer 128 Dim Delta Kernel linear+linear: 100%|██████████| 2000/2000 [04:17<00:00,  7.76it/s, loss=4.6250, acc=0.2714]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import random\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "from apex.normalization import FusedRMSNorm as RMSNorm\n",
    "from matplotlib import pyplot as plt\n",
    "# class RoundSTE(torch.autograd.Function):\n",
    "#     @staticmethod\n",
    "#     def forward(ctx, x):\n",
    "#         # 定义要靠近的整数值\n",
    "#         xdtype = x.dtype\n",
    "#         targets = torch.tensor([-1., 0., 1., 2.], device=x.device)\n",
    "        \n",
    "#         # 扩展维度以便广播计算\n",
    "#         distances = torch.abs(x.unsqueeze(-1) - targets)\n",
    "#         min_distances, nearest_indices = distances.min(dim=-1)\n",
    "        \n",
    "#         # 获取最近的整数值\n",
    "#         nearest_ints = targets[nearest_indices]\n",
    "        \n",
    "#         # 判断是否在 0.1 范围内\n",
    "#         mask = (min_distances <= 0.1).float()\n",
    "        \n",
    "#         # 如果在范围内，则用整数替代；否则保留原始值\n",
    "#         output = mask * nearest_ints + (1 - mask) * x\n",
    "        \n",
    "#         # 保存mask用于后向传播\n",
    "#         ctx.save_for_backward(mask)\n",
    "#         ctx.mark_non_differentiable(mask)\n",
    "        \n",
    "#         return output.to(xdtype)\n",
    "    \n",
    "#     @staticmethod\n",
    "#     def backward(ctx, grad_output):\n",
    "#         # 获取前向传播中保存的mask\n",
    "#         mask, = ctx.saved_tensors\n",
    "        \n",
    "#         # 对于被取整的部分，其梯度应该是0，因为它们被视为常量；\n",
    "#         # 对于未被取整的部分，梯度应该正常传递。\n",
    "#         grad_input = grad_output * (1 - mask)\n",
    "        \n",
    "#         return grad_input\n",
    "\n",
    "\n",
    "def optimize_vectors(num_vectors=32, dim=16, threshold=0.12, max_iter=5000, learning_rate=0.01):\n",
    "    # 检查是否有可用的 GPU\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    print(f\"Using device: {device}\")\n",
    "\n",
    "    # 随机初始化单位向量，并移动到设备上（CPU/GPU）\n",
    "    vectors = torch.randn(num_vectors, dim, device=device)\n",
    "    vectors /= vectors.norm(dim=1, keepdim=True)\n",
    "\n",
    "    # 设置为可求导变量\n",
    "    vectors.requires_grad_()\n",
    "\n",
    "    optimizer = torch.optim.Adam([vectors], lr=learning_rate)\n",
    "\n",
    "    for iteration in range(max_iter):\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 计算两两之间的内积矩阵，排除对角线（自己和自己）\n",
    "        dot_products = torch.mm(vectors, vectors.T) - torch.eye(num_vectors, device=device)\n",
    "\n",
    "        # 计算超过阈值的部分\n",
    "        excessive = dot_products.abs() - threshold\n",
    "        mask = (excessive > 0).float()\n",
    "\n",
    "        # 构造损失函数：只考虑超过阈值的项的平方\n",
    "        loss = (mask * excessive).pow(2).sum()\n",
    "\n",
    "        if loss.item() == 0:\n",
    "            print(f\"Converged at iteration {iteration}\")\n",
    "            break\n",
    "\n",
    "        # 反向传播与优化\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # 归一化向量\n",
    "        with torch.no_grad():\n",
    "            vectors /= vectors.norm(dim=1, keepdim=True)\n",
    "\n",
    "        if iteration % 100 == 0 or iteration == max_iter - 1:\n",
    "            print(f\"Iteration {iteration}, Current Loss: {loss.item():.6f}\")\n",
    "            if loss.item() < 0.001:\n",
    "                print(\"Loss is small enough. Breaking early.\")\n",
    "                break\n",
    "\n",
    "    # 返回最终结果作为 NumPy 数组（也可以保留为 Tensor）\n",
    "    return vectors.detach().cpu().numpy()\n",
    "\n",
    "\n",
    "\n",
    "class RoundSTE(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, x):\n",
    "        # 定义要靠近的整数值\n",
    "        xdtype = x.dtype\n",
    "\n",
    "        \n",
    "        return x.round().to(xdtype)\n",
    "    \n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output):\n",
    "\n",
    "        \n",
    "        return grad_output #grad_input\n",
    "    \n",
    "# Rotary positional embedding with configurable base (theta)\n",
    "def rotary_embedding(x, seq_dim=-2, base=10000):\n",
    "    seq_len = x.shape[seq_dim]\n",
    "    dim = x.shape[-1]\n",
    "    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device).float() / dim))\n",
    "    positions = torch.arange(seq_len, device=x.device).float()\n",
    "    freqs = torch.einsum('n,d->nd', positions, inv_freq)\n",
    "    emb = torch.cat((freqs.sin(), freqs.cos()), dim=-1)\n",
    "    return (x * emb) + (rotate_half(x) * emb)\n",
    "\n",
    "# helper to rotate half the dims\n",
    "def rotate_half(x):\n",
    "    x1, x2 = x.chunk(2, dim=-1)\n",
    "    return torch.cat((-x2, x1), dim=-1)\n",
    "\n",
    "class Embedding(nn.Module):\n",
    "    def __init__(self, vocab_size, dim):\n",
    "        super().__init__()\n",
    "        self.wte = nn.Embedding(vocab_size, dim)\n",
    "    def forward(self, x):\n",
    "        return self.wte(x)\n",
    "\n",
    "class Unembedding(nn.Module):\n",
    "    def __init__(self, dim, vocab_size):\n",
    "        super().__init__()\n",
    "        self.lm_head = nn.Linear(dim, vocab_size, bias=False)\n",
    "    def forward(self, x):\n",
    "        return self.lm_head(x)\n",
    "\n",
    "# Generalized Attention supporting MQA, GQA, Full MHA\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, dim, num_q_heads=8, num_kv_heads=None, rotary_base=10000):\n",
    "        super().__init__()\n",
    "        self.num_q_heads = num_q_heads\n",
    "        self.num_kv_heads = num_kv_heads or num_q_heads  # if None, full MHA\n",
    "        self.rotary_base = rotary_base\n",
    "        self.head_dim = dim // self.num_q_heads\n",
    "        self.q_proj = nn.Linear(dim, dim, bias=True)\n",
    "        self.k_proj = nn.Linear(dim, self.head_dim * num_kv_heads)\n",
    "        self.v_proj = nn.Linear(dim, self.head_dim * num_kv_heads)\n",
    "        self.out = nn.Linear(dim, dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, t, d = x.size()\n",
    "        # project\n",
    "        q = self.q_proj(x).view(b, t, self.num_q_heads, self.head_dim).transpose(1, 2)\n",
    "        k = self.k_proj(x).view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2)\n",
    "        v = self.v_proj(x).view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2)\n",
    "        # rotary on q,k\n",
    "        if self.rotary_base is not None:\n",
    "            q = rotary_embedding(q, base=self.rotary_base)\n",
    "            k = rotary_embedding(k, base=self.rotary_base)\n",
    "        # compute attention scores: broadcast kv heads to q heads if needed\n",
    "        k_expanded = k.unsqueeze(2).expand(-1, -1, self.num_q_heads // self.num_kv_heads, -1, -1).reshape(b, -1, t, self.head_dim)\n",
    "        v_expanded = v.unsqueeze(2).expand(-1, -1, self.num_q_heads // self.num_kv_heads, -1, -1).reshape(b, -1, t, self.head_dim)\n",
    "        # q: b,t,Qh,hd, k_expanded: b,t,Qh,hd\n",
    "        scores = torch.einsum('bhtd,bhsd->bhts', q, k_expanded) / math.sqrt(d // self.num_q_heads)\n",
    "        attn_mask = torch.tril(torch.ones(t, t)).unsqueeze(0).unsqueeze(1).to(x.device)\n",
    "        scores = scores.masked_fill(attn_mask == 0, float('-inf'))\n",
    "        attn = torch.softmax(scores, dim=-1).to(x.dtype)\n",
    "        out = torch.einsum('bhts,bhsd->bthd', attn, v_expanded)\n",
    "        out = out.reshape(b, t, d)\n",
    "        return self.out(out)\n",
    "\n",
    "class DeltaAttention(MultiHeadAttention):\n",
    "    def __init__(self, dim, num_q_heads=8, num_kv_heads=None, rotary_base=10000, kernel_1='linear', kernel_2='linear',vectors=None):\n",
    "        super().__init__(dim, num_q_heads, num_kv_heads, rotary_base)\n",
    "        self.alpha = nn.Parameter(torch.ones(1))\n",
    "        self.beta = nn.Parameter(torch.ones(1))\n",
    "        self.kernel_1 = kernel_1\n",
    "        self.kernel_2 = kernel_2\n",
    "\n",
    "        self.register_buffer('query',vectors[0])\n",
    "        self.register_buffer('key',vectors.reshape(1,1,vectors.shape[0],vectors.shape[1]))\n",
    "        self.register_buffer('value',vectors.reshape(1,1,vectors.shape[0],vectors.shape[1]))\n",
    "\n",
    "        self.init_length = vectors.shape[0]\n",
    "\n",
    "        # 手动初始化权重和偏置\n",
    "        nn.init.zeros_(self.q_proj.weight)\n",
    "        nn.init.zeros_(self.v_proj.weight)\n",
    "        nn.init.eye_(self.k_proj.weight)\n",
    "\n",
    "        self.q_proj.bias.data = vectors[0].clone().detach()\n",
    "        if self.k_proj.bias is not None:\n",
    "            nn.init.constant_(self.k_proj.bias, 0)\n",
    "        if self.v_proj.bias is not None:\n",
    "            nn.init.constant_(self.k_proj.bias, 0)\n",
    "    def forward(self, x,p_length=None):\n",
    "        b, t, d = x.size()\n",
    "        # project\n",
    "        ##data-independent query  + 恒等的key_proj + 0 proj的value  \n",
    "        q =  self.q_proj(x).view(b, t, self.num_q_heads, self.head_dim).transpose(1, 2) #self.q_proj(x).view(b, t, self.num_q_heads, self.head_dim).transpose(1, 2) #self.query + 0*self.q_proj(x).view(b, t, self.num_q_heads, self.head_dim).transpose(1, 2) #\n",
    "        k =    x.view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2)   # self.k_proj(x).view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2) # \n",
    "        v =   0*x.view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2) # v proj is 0.self.v_proj(x).view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2) #\n",
    "\n",
    "        # rotary on q,k\n",
    "        # if self.rotary_base is not None:\n",
    "        #     q = rotary_embedding(q, base=self.rotary_base)\n",
    "        #     k = rotary_embedding(k, base=self.rotary_base)\n",
    "        # compute attention scores: broadcast kv heads to q heads if needed\n",
    "        k_expanded = k.unsqueeze(2).expand(-1, -1, self.num_q_heads // self.num_kv_heads, -1, -1).reshape(b, -1, t, self.head_dim)\n",
    "        v_expanded = v.unsqueeze(2).expand(-1, -1, self.num_q_heads // self.num_kv_heads, -1, -1).reshape(b, -1, t, self.head_dim)\n",
    "\n",
    "        # print(self.k_expanded.shape,self.)\n",
    "        k_expanded = torch.cat((self.key.expand(b,-1,-1,-1),k_expanded),dim=2)\n",
    "        v_expanded = torch.cat((self.value.expand(b,-1,-1,-1),v_expanded),dim=2)\n",
    "        # A = torch.einsum(\"b h s d, b h t d-> b h s t\", k_expanded, k_expanded) / math.sqrt(d // self.num_q_heads)\n",
    "        A = torch.einsum(\"b h s d, b h t d-> b h s t\", k_expanded, k_expanded)\n",
    "        if self.kernel_1 == 'linear':\n",
    "            A = A.float()\n",
    "        elif self.kernel_1 == 'round':\n",
    "            A = (RoundSTE.apply(A)).float()\n",
    "        elif  self.kernel_1 == 'exp':\n",
    "            A = torch.exp(A).float()\n",
    "        elif self.kernel_1 == 'softmax':\n",
    "            A = torch.softmax(A, dim=-1).float()\n",
    "        elif self.kernel_1 == 'relu':\n",
    "            A = torch.relu(A).float()\n",
    "        # print(A[0][0].shape)\n",
    "        # print(A[0][0])\n",
    "        # assert False\n",
    "        u = torch.linalg.solve_triangular(self.beta * A, self.alpha * v_expanded, upper=False, unitriangular=True).to(k_expanded)\n",
    "\n",
    "        k_expanded = k_expanded[:,:,self.init_length:,:] \n",
    "        u = u[:,:,self.init_length:,:]\n",
    "        \n",
    "        # scores = torch.einsum('bhtd,bhsd->bhts', q, k_expanded) / math.sqrt(d // self.num_q_heads)\n",
    "        scores = torch.einsum('bhtd,bhsd->bhts', q, k_expanded)\n",
    "        attn_mask = torch.tril(torch.ones(t, t)).unsqueeze(0).unsqueeze(1).to(x.device)\n",
    "\n",
    "        if self.kernel_2 == 'linear':\n",
    "            attn = scores.masked_fill(attn_mask == 0, 0)\n",
    "            out = torch.einsum('bhts,bhsd->bthd', attn, u)\n",
    "        elif self.kernel_2 == 'round':\n",
    "            attn = (RoundSTE.apply(scores)).masked_fill(attn_mask == 0, 0)\n",
    "            out = torch.einsum('bhts,bhsd->bthd', attn, u)\n",
    "        elif self.kernel_2 == 'softmax':\n",
    "            scores = scores.masked_fill(attn_mask == 0, float('-inf'))\n",
    "            attn = torch.softmax(scores, dim=-1).to(x.dtype)\n",
    "            out = torch.einsum('bhts,bhsd->bthd', attn, u)\n",
    "        elif self.kernel_2 =='relu':\n",
    "            attn = torch.relu(scores).masked_fill(attn_mask == 0, 0)\n",
    "            out = torch.einsum('bhts,bhsd->bthd', attn, u)\n",
    "        out = out.reshape(b, t, d)\n",
    "        # return out\n",
    "        return self.out(out)\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, dim, num_q_heads=8, num_kv_heads=None, rotary_base=10000, kernel_1=None, kernel_2=None,vectors=None):\n",
    "        super().__init__()\n",
    "        self.ln1 = nn.LayerNorm(dim)\n",
    "        if kernel_1 is None:\n",
    "            self.attn = MultiHeadAttention(dim, num_q_heads, num_kv_heads, rotary_base)\n",
    "        else:\n",
    "            self.attn = DeltaAttention(dim, num_q_heads, num_kv_heads, rotary_base, kernel_1, kernel_2,vectors=vectors)\n",
    "\n",
    "    def forward(self, x,p_length):\n",
    "        return self.attn(x,p_length)\n",
    "        x = x + self.attn(x,p_length)\n",
    "        return x\n",
    "\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, vocab_size, dim, n_layers=1,\n",
    "                 num_heads=8, num_kv_heads=None, rotary_base=10000, kernel_1=None, kernel_2=None,vectors = None):\n",
    "        super().__init__()\n",
    "\n",
    "        assert vocab_size == (vectors.shape[0] * (vectors.shape[0]-1)//2)\n",
    "        assert vectors.shape[1] == dim\n",
    "\n",
    "        emb =  [vectors[i] - vectors[j] for i in range(vectors.shape[0]) for j in range(i+1, vectors.shape[0])]\n",
    "\n",
    "        # 构建固定嵌入矩阵\n",
    "        emb_list = [vectors[i] - vectors[j] for i in range(vectors.shape[0]) \n",
    "                                          for j in range(i+1, vectors.shape[0])]\n",
    "        emb = torch.stack(emb_list)  # shape: [vocab_size, dim]\n",
    "        \n",
    "        # 注册为 buffer，使其成为模型的一部分但不会被优化\n",
    "        self.register_buffer('emb', emb)\n",
    "\n",
    "        # self.emb = Embedding(vocab_size, dim)\n",
    "        self.blocks = nn.ModuleList([\n",
    "            TransformerBlock(dim, num_heads, num_kv_heads, rotary_base, kernel_1, kernel_2,vectors)\n",
    "            for _ in range(n_layers)\n",
    "        ])\n",
    "        self.unemb = Unembedding(dim, vocab_size)\n",
    "        self.p_length = int(math.sqrt(vocab_size*2)) + 1\n",
    "        # self.prompt = nn.Parameter(0.02*torch.randn(1, self.p_length , dim)) #%%用来写入初始n_elements\n",
    "        # self.norm1 = RMSNorm(dim)\n",
    "        self.norm2 = RMSNorm(dim)\n",
    "    def forward(self, x):\n",
    "        # x = self.emb(x) \n",
    "        x = self.emb[x]\n",
    "        # x = torch.cat([self.prompt.expand(x.shape[0],-1,-1),x],dim=1)\n",
    "        # x = self.norm1(x)\n",
    "        for block in self.blocks:\n",
    "            x = x + block(x,self.p_length)\n",
    "        # x = x[:,self.p_length:]\n",
    "        x = self.norm2(x)\n",
    "        return self.unemb(x)\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 512\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "# print(swap_pairs )\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "        self.data = generate_S5_data(self.k, self.num_samples)\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \n",
    "        input_ids, labels = self.data[idx]\n",
    "        if idx == 0:\n",
    "            self.data = generate_S5_data(self.k, self.num_samples)\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def train_model(model, dataloader, run_name=\"run\", log_dict=None,k=16,bs=32):\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 2000\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-2, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "    loss = nn.CrossEntropyLoss()\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "            # print(input_ids, labels)\n",
    "            logits = model(input_ids)\n",
    "            loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.shape[-1]), labels.view(-1))\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            progress_bar.set_postfix({\"loss\": f\"{loss.item():.4f}\", \"acc\": f\"{accuracy:.4f}\"})\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "        # if accuracy>0.999:\n",
    "        #     k*=2\n",
    "        #     print(f'step {step}, accuracy>0.99, length to {k}')\n",
    "        #     dataset = S5Dataset(k=k, num_samples=bs)\n",
    "        #     dataloader = DataLoader(dataset, batch_size=bs, shuffle=False)\n",
    "        #     if k>4096:\n",
    "        #         break\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "\n",
    "max_dot = 1\n",
    "dim = 128\n",
    "k = 1024\n",
    "while max_dot>0.125:\n",
    "    num_vectors  = n_elements\n",
    "    # 测试函数\n",
    "    vectors = optimize_vectors(num_vectors,dim=dim,learning_rate=0.01)\n",
    "\n",
    "    # 验证结果\n",
    "    dot_products = np.dot(vectors, vectors.T) - np.eye(num_vectors)\n",
    "    max_dot = np.max(np.abs(dot_products))\n",
    "\n",
    "print(f\"find {num_vectors} vectors in dim {dim} with max_abs_dot = {max_dot}\")\n",
    "vectors = torch.Tensor(vectors)\n",
    "print(vectors.shape)\n",
    "\n",
    "\n",
    "bs = 32\n",
    "\n",
    "dataset = S5Dataset(k=k, num_samples=bs)\n",
    "dataloader = DataLoader(dataset, batch_size=bs, shuffle=False)\n",
    "n_layers = 1\n",
    "vocab_size = math.comb(n_elements, 2)\n",
    "\n",
    "num_heads = 1\n",
    "num_kv_heads = 1\n",
    "# model_base = Transformer(\n",
    "#     vocab_size=vocab_size,\n",
    "#     dim=dim,\n",
    "#     n_layers=n_layers,\n",
    "#     num_heads=num_heads,\n",
    "#     num_kv_heads=num_kv_heads,\n",
    "#     rotary_base=None,\n",
    "# )\n",
    "kernel_cmb = [\n",
    "    ['round', 'round'],\n",
    "    # ['linear', 'softmax'],\n",
    "    ['linear','linear'],\n",
    "    # ['exp', 'softmax'],\n",
    "    # ['round', 'softmax'],\n",
    "]\n",
    "delta_models = []\n",
    "for kernels in kernel_cmb:\n",
    "\n",
    "    delta_models.append(Transformer(\n",
    "        vocab_size=vocab_size,\n",
    "        dim=dim,\n",
    "        n_layers=n_layers,\n",
    "        num_heads=num_heads,\n",
    "        num_kv_heads=num_kv_heads,\n",
    "        rotary_base=None,\n",
    "        kernel_1=kernels[0],\n",
    "        kernel_2=kernels[1],\n",
    "        vectors = vectors,\n",
    "    ))\n",
    "\n",
    "log_dict = {}\n",
    "# train_model(model_base, dataloader, run_name=f'{n_layers} Layer {dim} Dim Transformer', log_dict=log_dict)\n",
    "for i, kernels in enumerate(kernel_cmb):\n",
    "    train_model(delta_models[i], dataloader, run_name=f'{n_layers} Layer {dim} Dim Delta Kernel {kernels[0]}+{kernels[1]}', log_dict=log_dict,k=k,bs=bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 100592 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(12, 6))\n",
    "for key in log_dict:\n",
    "    accs = log_dict[key]['acc']\n",
    "    steps = np.arange(len(accs))\n",
    "    if 'round' in key:\n",
    "        key = 'round'\n",
    "    else:\n",
    "        key = 'linear'\n",
    "    plt.plot(steps, accs, label=key)\n",
    "\n",
    "# plt.title(\"Accuracy vs Training Steps\")\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "fileId": "d745737a-a4c3-4b71-b34e-7d63fae1737c",
  "filePath": "/mnt/bn/seed-aws-va/zhongshu/research/deltaformer/swap.ipynb",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
