{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10 | Loss: 9.3827\n",
      "Epoch 2/10 | Loss: 9.3907\n",
      "Epoch 3/10 | Loss: 9.3653\n",
      "Epoch 4/10 | Loss: 9.3593\n",
      "Epoch 5/10 | Loss: 9.3731\n",
      "Epoch 6/10 | Loss: 9.3543\n",
      "Epoch 7/10 | Loss: 9.3582\n",
      "Epoch 8/10 | Loss: 9.3560\n",
      "Epoch 9/10 | Loss: 9.3424\n",
      "Epoch 10/10 | Loss: 9.3296\n"
     ]
    }
   ],
   "source": [
    "import torch, math\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Positional Encoding\n",
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=5000):\n",
    "        super().__init__()\n",
    "        pos = torch.arange(max_len).unsqueeze(1)\n",
    "        div = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        pe[:, 0::2], pe[:, 1::2] = torch.sin(pos * div), torch.cos(pos * div)\n",
    "        self.pe = pe.unsqueeze(0).to(device)\n",
    "\n",
    "    def forward(self, x): return x + self.pe[:, :x.size(1)]\n",
    "\n",
    "# Multi-Head Self-Attention\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model, heads):\n",
    "        super().__init__()\n",
    "        assert d_model % heads == 0\n",
    "        self.head_dim = d_model // heads\n",
    "        self.qkv = nn.Linear(d_model, d_model * 3)\n",
    "        self.out = nn.Linear(d_model, d_model)\n",
    "        self.scale = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(device)\n",
    "        self.heads = heads\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, L, D = x.size()\n",
    "        qkv = self.qkv(x).view(B, L, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4)\n",
    "        q, k, v = qkv\n",
    "        scores = (q @ k.transpose(-2, -1)) / self.scale\n",
    "        attn = torch.softmax(scores, dim=-1)\n",
    "        out = (attn @ v).permute(0, 2, 1, 3).reshape(B, L, D)\n",
    "        return self.out(out)\n",
    "\n",
    "# Feedforward Layer\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, d_model, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(d_model, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, d_model)\n",
    "        )\n",
    "\n",
    "    def forward(self, x): return self.net(x)\n",
    "\n",
    "# Transformer Encoder Layer\n",
    "class TransformerEncoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, heads, hidden_dim, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.attn = MultiHeadAttention(d_model, heads)\n",
    "        self.ffn = FeedForward(d_model, hidden_dim)\n",
    "        self.norm1 = nn.LayerNorm(d_model)\n",
    "        self.norm2 = nn.LayerNorm(d_model)\n",
    "        self.drop = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.norm1(x + self.drop(self.attn(x)))\n",
    "        return self.norm2(x + self.drop(self.ffn(x)))\n",
    "\n",
    "# Transformer Encoder\n",
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, d_model, heads, layers, hidden_dim, vocab_size, max_len):\n",
    "        super().__init__()\n",
    "        self.embed = nn.Embedding(vocab_size, d_model)\n",
    "        self.pos_enc = PositionalEncoding(d_model, max_len)\n",
    "        self.enc_layers = nn.ModuleList([TransformerEncoderLayer(d_model, heads, hidden_dim) for _ in range(layers)])\n",
    "        self.out = nn.Linear(d_model, vocab_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pos_enc(self.embed(x))\n",
    "        for layer in self.enc_layers: x = layer(x)\n",
    "        return self.out(x)\n",
    "\n",
    "# Hyperparameters & Training\n",
    "d_model, heads, layers, hidden_dim, vocab_size, max_len = 128, 8, 6, 512, 10000, 100\n",
    "model = TransformerEncoder(d_model, heads, layers, hidden_dim, vocab_size, max_len).to(device)\n",
    "criterion, optimizer = nn.CrossEntropyLoss(), optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "# Dummy training loop\n",
    "def train_model(epochs=10):\n",
    "    for epoch in range(epochs):\n",
    "        x = torch.randint(0, vocab_size, (32, max_len)).to(device)\n",
    "        y = torch.randint(0, vocab_size, (32, max_len)).to(device)\n",
    "        optimizer.zero_grad()\n",
    "        loss = criterion(model(x).view(-1, vocab_size), y.view(-1))\n",
    "        loss.backward(); optimizer.step()\n",
    "        print(f\"Epoch {epoch+1}/{epochs} | Loss: {loss.item():.4f}\")\n",
    "\n",
    "train_model()\n"
   ]
  }
 ],
 "metadata": {
  "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}