add files

README.md

@@ -1,13 +0,0 @@
----
-title: Chatbot
-emoji: 🚀
-colorFrom: purple
-colorTo: gray
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
-short_description: they call me sam altman
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,95 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from components.model import GPTModel
+from components.tokenizer import encode, decode, tokenizer
+
+
+# -----------------------------
+# Load model & configuration
+# -----------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Hyperparameters should match training
+block_size = 128
+n_layers = 16
+n_heads = 8
+dropout_p = 0.1
+n_embedding = 256
+
+# initialize model and load weights
+vocab_size = tokenizer.n_vocab
+model = GPTModel(vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size).to(
+    device
+)
+model.load_state_dict(torch.load("checkpoints/gpt_model-1.pth", map_location=device))
+model.eval()
+
+
+# -----------------------------
+# Generation function
+# -----------------------------
+@torch.no_grad()
+def generate_text(prompt, max_new_tokens=200, temperature=1.0, top_k=50):
+    model.eval()
+
+    # Wrap message in [INST] and [/INST]
+    wrapped_prompt = f"[INST] {prompt.strip()} [/INST]"
+    tokens = (
+        torch.tensor(encode(wrapped_prompt), dtype=torch.long).unsqueeze(0).to(device)
+    )
+
+    inst_token_id = encode("[INST]")[0]
+
+    for _ in range(max_new_tokens):
+        input_tokens = tokens[:, -block_size:]
+        logits = model(input_tokens)
+        logits = logits[:, -1, :] / temperature
+
+        if top_k is not None:
+            values, indices = torch.topk(logits, top_k)
+            logits[logits < values[:, [-1]]] = -float("Inf")
+
+        probs = F.softmax(logits, dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1)
+
+        # Stop generation if [INST] appears again (do not include it)
+        if next_token.item() == inst_token_id:
+            break
+
+        tokens = torch.cat((tokens, next_token), dim=1)
+
+    return decode(tokens[0].tolist())[len(wrapped_prompt) :]
+
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+def chat(prompt, max_tokens, temperature, top_k):
+    response = generate_text(prompt, max_tokens, temperature, top_k)
+    return response
+
+
+with gr.Blocks(title="TinyChat GPT Model") as demo:
+    gr.Markdown("## cute lil chatbot")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            prompt = gr.Textbox(
+                label="Prompt", placeholder="Type your message here...", lines=4
+            )
+            max_tokens = gr.Slider(10, 500, value=200, step=10, label="Max New Tokens")
+            temperature = gr.Slider(0.2, 1.5, value=1.0, step=0.1, label="Temperature")
+            top_k = gr.Slider(10, 200, value=50, step=10, label="Top‑K Sampling")
+            submit = gr.Button("Generate")
+
+        with gr.Column(scale=3):
+            output = gr.Textbox(label="Generated Response", lines=15)
+
+    submit.click(chat, inputs=[prompt, max_tokens, temperature, top_k], outputs=output)
+
+# -----------------------------
+# Launch app
+# -----------------------------
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)

components/dataset.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math, time, os
+from torch.utils.data import Dataset, DataLoader
+import tiktoken
+
+# from torch.cuda.amp import autocast, GradScaler
+from torch.amp.autocast_mode import autocast
+from torch.amp.grad_scaler import GradScaler
+from tqdm import tqdm
+
+from datasets import load_dataset
+from components.model import GPTModel
+from components.tokenizer import encode, decode, tokenizer
+
+
+def decode(tokens):
+    return tokenizer.decode(tokens)
+
+class TextDataset(Dataset):
+    def __init__(self, hf_dataset, block_size):
+        self.dataset = hf_dataset
+        # self.tokenizer = tokenizer
+        self.block_size = block_size
+
+    def __len__(self):
+        return len(self.dataset["train"])
+
+    def __getitem__(self, idx):
+        # Start with a random index sample
+        rand_idx = torch.randint(0, len(self.dataset["train"]), (1,)).item()
+        text = self.dataset["train"][rand_idx]["text"]
+        tokens = encode(text)
+
+        # Keep appending more samples if too short
+        while len(tokens) < self.block_size + 1:
+            next_idx = torch.randint(0, len(self.dataset["train"]), (1,)).item()
+            next_text = self.dataset["train"][next_idx]["text"]
+            tokens.extend(encode(" " + next_text))
+            # Prevent runaway growth
+            if len(tokens) > self.block_size * 2:
+                break
+
+        # Truncate to block_size + 1
+        tokens = torch.tensor(tokens[: self.block_size + 1])
+
+        x = tokens[: self.block_size]
+        y = tokens[1 : self.block_size + 1]
+        return x.long(), y.long()

components/model.py

@@ -0,0 +1,153 @@
+
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math, time, os
+from torch.utils.data import Dataset, DataLoader
+import tiktoken
+
+# from torch.cuda.amp import autocast, GradScaler
+from torch.amp.autocast_mode import autocast
+from torch.amp.grad_scaler import GradScaler
+from tqdm import tqdm
+
+from datasets import load_dataset
+
+
+# class GPTModel(nn.Module):
+#     def __init__(
+#         self, vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size
+#     ):
+#         super(GPTModel, self).__init__()
+#         self.token_embedding = nn.Embedding(vocab_size, n_embedding)
+#         self.position_embedding = nn.Embedding(block_size, n_embedding)
+#         self.layers = nn.ModuleList(
+#             [
+#                 nn.TransformerEncoderLayer(
+#                     d_model=n_embedding, nhead=n_heads, dropout=dropout_p
+#                 )
+#                 for _ in range(n_layers)
+#             ]
+#         )
+#         self.ln_f = nn.LayerNorm(n_embedding)
+#         self.head = nn.Linear(n_embedding, vocab_size)
+#         self.dropout = nn.Dropout(dropout_p)
+#         self.block_size = block_size
+
+#     def forward(self, x):
+#         bsz, seq_len = x.size()
+#         positions = (
+#             torch.arange(0, seq_len, device=x.device).unsqueeze(0).expand(bsz, seq_len)
+#         )
+#         x = self.token_embedding(x) + self.position_embedding(positions)
+#         x = self.dropout(x)
+
+#         for layer in self.layers:
+#             x = layer(x)
+
+#         x = self.ln_f(x)
+#         logits = self.head(x)
+#         return logits
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+# ... existing imports ...
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, n_embedding, n_heads, dropout_p):
+        super().__init__()
+        assert n_embedding % n_heads == 0
+        self.n_heads = n_heads
+        self.head_dim = n_embedding // n_heads
+
+        self.q_proj = nn.Linear(n_embedding, n_embedding)
+        self.k_proj = nn.Linear(n_embedding, n_embedding)
+        self.v_proj = nn.Linear(n_embedding, n_embedding)
+        self.out_proj = nn.Linear(n_embedding, n_embedding)
+        self.dropout = nn.Dropout(dropout_p)
+
+    def forward(self, x, attn_mask=None):
+        B, T, C = x.shape # batch size, seq length, embedding dim
+
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        # c, embd dim split into n_heads x head_dim
+
+        # built-in scaled dot product attention for efficiency
+        attn_out = F.scaled_dot_product_attention(
+            q, k, v, 
+            attn_mask=attn_mask,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            is_causal=True,
+        )
+
+        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.out_proj(attn_out)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, n_embedding, dropout_p):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(n_embedding, 4 * n_embedding),
+            nn.GELU(),
+            nn.Linear(4 * n_embedding, n_embedding),
+            nn.Dropout(dropout_p),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, n_embedding, n_heads, dropout_p):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(n_embedding)
+        self.ln2 = nn.LayerNorm(n_embedding)
+        self.attn = MultiHeadAttention(n_embedding, n_heads, dropout_p)
+        self.ff = FeedForward(n_embedding, dropout_p)
+
+    def forward(self, x, attn_mask=None):
+        x = x + self.attn(self.ln1(x), attn_mask)
+        x = x + self.ff(self.ln2(x))
+        return x
+
+
+class GPTModel(nn.Module):
+    def __init__(self, vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size):
+        super().__init__()
+        self.token_embed = nn.Embedding(vocab_size, n_embedding)
+        self.pos_embed = nn.Embedding(block_size, n_embedding)
+        self.blocks = nn.ModuleList([
+            TransformerBlock(n_embedding, n_heads, dropout_p) 
+            for _ in range(n_layers)
+        ])
+        self.ln_f = nn.LayerNorm(n_embedding)
+        self.head = nn.Linear(n_embedding, vocab_size, bias=False)
+        self.dropout = nn.Dropout(dropout_p)
+        self.block_size = block_size
+
+    def forward(self, idx):
+        B, T = idx.shape
+        assert T <= self.block_size, "Sequence exceeds block size."
+
+        pos = torch.arange(0, T, device=idx.device).unsqueeze(0)
+        x = self.token_embed(idx) + self.pos_embed(pos)
+        x = self.dropout(x)
+
+        # Causal mask for decoder: prevent attending to future tokens
+        attn_mask = torch.ones(T, T, device=idx.device, dtype=torch.bool).tril()
+
+        for block in self.blocks:
+            x = block(x, attn_mask)
+
+        x = self.ln_f(x)
+        logits = self.head(x)
+        return logits

components/tokenizer.py

@@ -0,0 +1,41 @@
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math, time, os
+from torch.utils.data import Dataset, DataLoader
+import tiktoken
+
+# from torch.cuda.amp import autocast, GradScaler
+from torch.amp.autocast_mode import autocast
+from torch.amp.grad_scaler import GradScaler
+from tqdm import tqdm
+
+from datasets import load_dataset
+from components.model import GPTModel
+
+
+tokenizer = tiktoken.get_encoding("gpt2")
+
+base_encoding = tiktoken.get_encoding("gpt2")
+
+special_tokens = {
+    "[INST]": base_encoding.n_vocab,  # next available token id
+    "[/INST]": base_encoding.n_vocab + 1,
+}
+
+# 3. Create a new encoding that merges GPT‑2’s tokens + your special tokens
+tokenizer = tiktoken.Encoding(
+    name="gpt2_with_inst",
+    pat_str=base_encoding._pat_str,
+    mergeable_ranks=base_encoding._mergeable_ranks,
+    special_tokens={**base_encoding._special_tokens, **special_tokens},
+)
+
+
+def encode(text):
+    return tokenizer.encode(text, allowed_special={"[INST]", "[/INST]"})
+
+
+def decode(tokens):
+    return tokenizer.decode(tokens)

old/rl_test.ipynb

@@ -0,0 +1,502 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "158eaa47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.nn import functional as F\n",
+    "import math, time, os\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "import tiktoken\n",
+    "# from torch.cuda.amp import autocast, GradScaler\n",
+    "from torch.amp.autocast_mode import autocast\n",
+    "from torch.amp.grad_scaler import GradScaler"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "97d9467e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/software/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Her campaign emailed a fundraising pitch Tuesday evening warning of the dangers of a Trump presidency and of complacency among Democrats.\n",
+      "{'text': \"Canonical, keeper of the Ubuntu Linux distribution, is a small company with big friends. The latest example: Dell, IBM and Intel each are taking new steps with Ubuntu. Here's the scoop.\"}\n"
+     ]
+    }
+   ],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "# dataset = load_dataset(\"wikimedia/wikipedia\", \"20231101.en\")\n",
+    "dataset = load_dataset(\"Bingsu/openwebtext_20p\")\n",
+    "ds = load_dataset(\"starhopp3r/TinyChat\", split=\"train\")\n",
+    "# This gives you cleaned, plain text articles1\n",
+    "print(dataset['train'][100]['text'][:500])  # Print the first 500 characters of the first article\n",
+    "print(dataset['train'][600000])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "81b98c54",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class TextDataset(Dataset):\n",
+    "    def __init__(self, hf_dataset, tokenizer, block_size):\n",
+    "        self.dataset = hf_dataset\n",
+    "        self.tokenizer = tokenizer\n",
+    "        self.block_size = block_size\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.dataset['train'])\n",
+    "\n",
+    "    # def __getitem__(self, idx):\n",
+    "    #     tokens = self.tokenizer.encode(self.dataset['train'][idx]['text'])\n",
+    "    #     if len(tokens) < self.block_size + 1:\n",
+    "    #         tokens = F.pad(torch.tensor(tokens), (0, self.block_size + 1 - len(tokens)), value=0)\n",
+    "    #     else:\n",
+    "    #         tokens = torch.tensor(tokens[: self.block_size + 1])\n",
+    "    #     x = tokens[: self.block_size]\n",
+    "    #     y = tokens[1 : self.block_size + 1]\n",
+    "    #     return x.long(), y.long()\n",
+    "    def __getitem__(self, idx):\n",
+    "        # choose a random index instead of using the passed idx\n",
+    "        rand_idx = torch.randint(0, len(self.dataset['train']), (1,)).item()\n",
+    "        tokens = self.tokenizer.encode(self.dataset['train'][rand_idx]['text'])\n",
+    "\n",
+    "        if len(tokens) < self.block_size + 1:\n",
+    "            tokens = F.pad(torch.tensor(tokens), (0, self.block_size + 1 - len(tokens)), value=0)\n",
+    "        else:\n",
+    "            tokens = torch.tensor(tokens[: self.block_size + 1])\n",
+    "\n",
+    "        x = tokens[: self.block_size]\n",
+    "        y = tokens[1 : self.block_size + 1]\n",
+    "        return x.long(), y.long()\n",
+    "\n",
+    "import torch\n",
+    "from torch.utils.data import Dataset\n",
+    "from datasets import load_dataset\n",
+    "import re\n",
+    "\n",
+    "class ChatDataset(Dataset):\n",
+    "    def __init__(self, tokenizer, split=\"train\", block_size=256, dataset_name=\"starhopp3r/TinyChat\"):\n",
+    "        \"\"\"\n",
+    "        Args:\n",
+    "            tokenizer: a tokenizer (e.g., tiktoken or Hugging Face tokenizer)\n",
+    "            split: dataset split (\"train\" etc)\n",
+    "            block_size: maximum sequence length\n",
+    "            dataset_name: path/name of the Hugging Face dataset\n",
+    "        \"\"\"\n",
+    "        self.dataset = load_dataset(dataset_name, split=split)\n",
+    "        self.tokenizer = tokenizer\n",
+    "        self.block_size = block_size\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.dataset)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        sample = self.dataset[idx]\n",
+    "        text = sample[\"text\"]\n",
+    "\n",
+    "        # --- split into prompt and response (TinyChat uses [INST] ... [/INST]) ---\n",
+    "        match = re.search(r\"\\[INST\\](.*?)\\[/INST\\](.*)\", text, re.DOTALL)\n",
+    "        if match:\n",
+    "            instruction = match.group(1).strip()\n",
+    "            response = match.group(2).strip()\n",
+    "        else:\n",
+    "            instruction = text.strip()\n",
+    "            response = \"\"\n",
+    "\n",
+    "        # Combine into a training sequence\n",
+    "        combined_text = f\"<inst> {instruction} </inst> {response}\"\n",
+    "\n",
+    "        # Tokenize (truncate/pad to block_size + 1)\n",
+    "        tokens = torch.tensor(self.tokenizer.encode(combined_text), dtype=torch.long)\n",
+    "        if len(tokens) < self.block_size + 1:\n",
+    "            pad_len = self.block_size + 1 - len(tokens)\n",
+    "            tokens = F.pad(tokens, (0, pad_len), value=0)\n",
+    "        else:\n",
+    "            tokens = tokens[: self.block_size + 1]\n",
+    "\n",
+    "        x = tokens[:-1]\n",
+    "        y = tokens[1:]\n",
+    "\n",
+    "        return x, y"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "599aa05a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#hyperparameters\n",
+    "train_model = True\n",
+    "compile_model = True\n",
+    "block_size = 256\n",
+    "n_layers = 32\n",
+    "n_heads = 16\n",
+    "dropout_p = 0.1\n",
+    "batch_size =16\n",
+    "learning_rate = 3e-4\n",
+    "n_embedding = 512\n",
+    "max_iters = 1000\n",
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "a69561e9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
+    "\n",
+    "train_dataset = TextDataset(dataset, tokenizer, block_size=block_size)\n",
+    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)\n",
+    "\n",
+    "chat_dataset = ChatDataset(tokenizer, split=\"train\", block_size=block_size)\n",
+    "chat_dataloader = DataLoader(chat_dataset, batch_size=batch_size, shuffle=True, drop_last=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "ea5598ea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class GPTModel(nn.Module):\n",
+    "    def __init__(self, vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size):\n",
+    "        super(GPTModel, self).__init__()\n",
+    "        self.token_embedding = nn.Embedding(vocab_size, n_embedding)\n",
+    "        self.position_embedding = nn.Embedding(block_size, n_embedding)\n",
+    "        self.layers = nn.ModuleList([\n",
+    "            nn.TransformerEncoderLayer(d_model=n_embedding, nhead=n_heads, dropout=dropout_p)\n",
+    "            for _ in range(n_layers)\n",
+    "        ])\n",
+    "        self.ln_f = nn.LayerNorm(n_embedding)\n",
+    "        self.head = nn.Linear(n_embedding, vocab_size)\n",
+    "        self.dropout = nn.Dropout(dropout_p)\n",
+    "        self.block_size = block_size\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        bsz, seq_len = x.size()\n",
+    "        positions = torch.arange(0, seq_len, device=x.device).unsqueeze(0).expand(bsz, seq_len)\n",
+    "        x = self.token_embedding(x) + self.position_embedding(positions)\n",
+    "        x = self.dropout(x)\n",
+    "\n",
+    "        for layer in self.layers:\n",
+    "            x = layer(x)\n",
+    "\n",
+    "        x = self.ln_f(x)\n",
+    "        logits = self.head(x)\n",
+    "        return logits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "6a1344ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define objects\n",
+    "vocab_size = tokenizer.n_vocab\n",
+    "\n",
+    "model = GPTModel(vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size).to(device)\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
+    "loss_fn = nn.CrossEntropyLoss()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "a0982489",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/software/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/__init__.py:1617: UserWarning: Please use the new API settings to control TF32 behavior, such as torch.backends.cudnn.conv.fp32_precision = 'tf32' or torch.backends.cuda.matmul.fp32_precision = 'ieee'. Old settings, e.g, torch.backends.cuda.matmul.allow_tf32 = True, torch.backends.cudnn.allow_tf32 = True, allowTF32CuDNN() and allowTF32CuBLAS() will be deprecated after Pytorch 2.9. Please see https://pytorch.org/docs/main/notes/cuda.html#tensorfloat-32-tf32-on-ampere-and-later-devices (Triggered internally at /pytorch/aten/src/ATen/Context.cpp:80.)\n",
+      "  _C._set_float32_matmul_precision(precision)\n",
+      "Training: 100%|████████████████████████████████████| 1000/1000 [05:06<00:00,  3.26it/s, loss=1.5960]\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "# training loop\n",
+    "scaler = GradScaler(device)\n",
+    "if train_model:\n",
+    "    if compile_model:\n",
+    "      compiled_model = torch.compile(model)\n",
+    "      torch.set_float32_matmul_precision('high')\n",
+    "    else:\n",
+    "      compiled_model = model\n",
+    "\n",
+    "    pbar = tqdm(range(max_iters), desc=\"Training\", ncols=100)\n",
+    "    data_iter = iter(train_dataloader)\n",
+    "    chat_data_iter = iter(chat_dataloader)\n",
+    "\n",
+    "    for count in pbar:\n",
+    "        try:\n",
+    "            if count %2 ==0:\n",
+    "              xb, yb = next(chat_data_iter)\n",
+    "            else:\n",
+    "              xb, yb = next(data_iter)\n",
+    "        except StopIteration:\n",
+    "            break  # dataloader exhausted before max_iters\n",
+    "        \n",
+    "        xb, yb = xb.to(device), yb.to(device)\n",
+    "        # logits = compiled_model(xb)\n",
+    "        # loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))\n",
+    "\n",
+    "        # optimizer.zero_grad()\n",
+    "        # loss.backward()\n",
+    "        # optimizer.step()\n",
+    "        with autocast(device, dtype=torch.float16):\n",
+    "            logits = compiled_model(xb)\n",
+    "            loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))\n",
+    "\n",
+    "        # backward pass with gradient scaling\n",
+    "        optimizer.zero_grad()\n",
+    "        scaler.scale(loss).backward()\n",
+    "        scaler.step(optimizer)\n",
+    "        scaler.update()\n",
+    "\n",
+    "        # update bar text dynamically\n",
+    "        pbar.set_postfix({\"loss\": f\"{loss.item():.4f}\"})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "6eb95580",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if train_model:\n",
+    "  torch.save(model.state_dict(), \"checkpoints/gpt_model-2.pth\")\n",
+    "else:\n",
+    "  model.load_state_dict(torch.load(\"checkpoints/gpt_model-2.pth\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "4371725d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "me when the .!!!!!! understand!!!!] cold! especially characters!! used soon!!!!! world! Exactly]-INST!!! choices! feel! spread!! a!! impact]inst saw them\n"
+     ]
+    }
+   ],
+   "source": [
+    "@torch.no_grad()\n",
+    "def generate_text(model, tokenizer, prompt, max_new_tokens, block_size, device):\n",
+    "    model.eval()\n",
+    "    # Encode the prompt text into token IDs\n",
+    "    tokens = torch.tensor(tokenizer.encode(prompt), dtype=torch.long).unsqueeze(0).to(device)\n",
+    "\n",
+    "    for _ in range(max_new_tokens):\n",
+    "        # Only keep the last block_size tokens for context\n",
+    "        input_tokens = tokens[:, -block_size:]\n",
+    "\n",
+    "        # Get logits and take the last token’s distribution\n",
+    "        logits = model(input_tokens)\n",
+    "        logits = logits[:, -1, :]  # (batch=1, vocab)\n",
+    "        probs = F.softmax(logits, dim=-1)\n",
+    "\n",
+    "        # Sample from the distribution\n",
+    "        next_token = torch.multinomial(probs, num_samples=1)\n",
+    "        tokens = torch.cat((tokens, next_token), dim=1)\n",
+    "\n",
+    "    # Decode back into text\n",
+    "    output_text = tokenizer.decode(tokens[0].tolist())\n",
+    "    return output_text\n",
+    "  \n",
+    "prompt = \"me when the \"\n",
+    "print(generate_text(model, tokenizer, prompt, max_new_tokens=50, block_size=block_size, device=device))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "d9c83f71",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 1/500 [00:01<12:16,  1.48s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Step 0: reward=0.00\n",
+      "Generated:\n",
+      "Hello:!!!!!!! that [ it everywhere [ </!!! impact!! not!!!!!!! past! un!,. to especially! explanation now! colorful, more!>!!!]! [/\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      " 10%|▉         | 49/500 [01:09<10:43,  1.43s/it]\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+      "\u001b[31mKeyboardInterrupt\u001b[39m                         Traceback (most recent call last)",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[11]\u001b[39m\u001b[32m, line 50\u001b[39m\n\u001b[32m     47\u001b[39m prompt = sample.get(\u001b[33m\"\u001b[39m\u001b[33mprompt\u001b[39m\u001b[33m\"\u001b[39m) \u001b[38;5;129;01mor\u001b[39;00m sample.get(\u001b[33m\"\u001b[39m\u001b[33minput\u001b[39m\u001b[33m\"\u001b[39m) \u001b[38;5;129;01mor\u001b[39;00m \u001b[33m\"\u001b[39m\u001b[33mHello:\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m     49\u001b[39m \u001b[38;5;66;03m# 2. Generate text and token logprobs\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m50\u001b[39m text, logprob_sum = \u001b[43mgenerate_with_logprobs\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtokenizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprompt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mblock_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmax_new_tokens\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     52\u001b[39m \u001b[38;5;66;03m# 3. Compute reward\u001b[39;00m\n\u001b[32m     53\u001b[39m r = torch.tensor(reward_fn(text), device=device)\n",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[11]\u001b[39m\u001b[32m, line 28\u001b[39m, in \u001b[36mgenerate_with_logprobs\u001b[39m\u001b[34m(model, tokenizer, prompt, block_size, max_new_tokens)\u001b[39m\n\u001b[32m     26\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(max_new_tokens):\n\u001b[32m     27\u001b[39m     input_tokens = tokens[:, -block_size:]\n\u001b[32m---> \u001b[39m\u001b[32m28\u001b[39m     logits = \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_tokens\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     29\u001b[39m     logits = logits[:, -\u001b[32m1\u001b[39m, :]  \u001b[38;5;66;03m# (1, vocab)\u001b[39;00m\n\u001b[32m     30\u001b[39m     probs = F.softmax(logits, dim=-\u001b[32m1\u001b[39m)\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1775\u001b[39m, in \u001b[36mModule._wrapped_call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1773\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._compiled_call_impl(*args, **kwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[32m   1774\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1775\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1786\u001b[39m, in \u001b[36mModule._call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1781\u001b[39m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[32m   1782\u001b[39m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[32m   1783\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m._backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_pre_hooks\n\u001b[32m   1784\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[32m   1785\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[32m-> \u001b[39m\u001b[32m1786\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1788\u001b[39m result = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m   1789\u001b[39m called_always_called_hooks = \u001b[38;5;28mset\u001b[39m()\n",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[6]\u001b[39m\u001b[32m, line 22\u001b[39m, in \u001b[36mGPTModel.forward\u001b[39m\u001b[34m(self, x)\u001b[39m\n\u001b[32m     19\u001b[39m x = \u001b[38;5;28mself\u001b[39m.dropout(x)\n\u001b[32m     21\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.layers:\n\u001b[32m---> \u001b[39m\u001b[32m22\u001b[39m     x = \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     24\u001b[39m x = \u001b[38;5;28mself\u001b[39m.ln_f(x)\n\u001b[32m     25\u001b[39m logits = \u001b[38;5;28mself\u001b[39m.head(x)\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1775\u001b[39m, in \u001b[36mModule._wrapped_call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1773\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._compiled_call_impl(*args, **kwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[32m   1774\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1775\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1786\u001b[39m, in \u001b[36mModule._call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1781\u001b[39m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[32m   1782\u001b[39m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[32m   1783\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m._backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_pre_hooks\n\u001b[32m   1784\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[32m   1785\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[32m-> \u001b[39m\u001b[32m1786\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1788\u001b[39m result = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m   1789\u001b[39m called_always_called_hooks = \u001b[38;5;28mset\u001b[39m()\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/transformer.py:935\u001b[39m, in \u001b[36mTransformerEncoderLayer.forward\u001b[39m\u001b[34m(self, src, src_mask, src_key_padding_mask, is_causal)\u001b[39m\n\u001b[32m    931\u001b[39m     x = x + \u001b[38;5;28mself\u001b[39m._ff_block(\u001b[38;5;28mself\u001b[39m.norm2(x))\n\u001b[32m    932\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m    933\u001b[39m     x = \u001b[38;5;28mself\u001b[39m.norm1(\n\u001b[32m    934\u001b[39m         x\n\u001b[32m--> \u001b[39m\u001b[32m935\u001b[39m         + \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_sa_block\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msrc_mask\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msrc_key_padding_mask\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mis_causal\u001b[49m\u001b[43m=\u001b[49m\u001b[43mis_causal\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    936\u001b[39m     )\n\u001b[32m    937\u001b[39m     x = \u001b[38;5;28mself\u001b[39m.norm2(x + \u001b[38;5;28mself\u001b[39m._ff_block(x))\n\u001b[32m    939\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m x\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/transformer.py:949\u001b[39m, in \u001b[36mTransformerEncoderLayer._sa_block\u001b[39m\u001b[34m(self, x, attn_mask, key_padding_mask, is_causal)\u001b[39m\n\u001b[32m    942\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_sa_block\u001b[39m(\n\u001b[32m    943\u001b[39m     \u001b[38;5;28mself\u001b[39m,\n\u001b[32m    944\u001b[39m     x: Tensor,\n\u001b[32m   (...)\u001b[39m\u001b[32m    947\u001b[39m     is_causal: \u001b[38;5;28mbool\u001b[39m = \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[32m    948\u001b[39m ) -> Tensor:\n\u001b[32m--> \u001b[39m\u001b[32m949\u001b[39m     x = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mself_attn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    950\u001b[39m \u001b[43m        \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    951\u001b[39m \u001b[43m        \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    952\u001b[39m \u001b[43m        \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    953\u001b[39m \u001b[43m        \u001b[49m\u001b[43mattn_mask\u001b[49m\u001b[43m=\u001b[49m\u001b[43mattn_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    954\u001b[39m \u001b[43m        \u001b[49m\u001b[43mkey_padding_mask\u001b[49m\u001b[43m=\u001b[49m\u001b[43mkey_padding_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    955\u001b[39m \u001b[43m        \u001b[49m\u001b[43mneed_weights\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m    956\u001b[39m \u001b[43m        \u001b[49m\u001b[43mis_causal\u001b[49m\u001b[43m=\u001b[49m\u001b[43mis_causal\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    957\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m[\u001b[32m0\u001b[39m]\n\u001b[32m    958\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m.dropout1(x)\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1775\u001b[39m, in \u001b[36mModule._wrapped_call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1773\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._compiled_call_impl(*args, **kwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[32m   1774\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1775\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1786\u001b[39m, in \u001b[36mModule._call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1781\u001b[39m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[32m   1782\u001b[39m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[32m   1783\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m._backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_pre_hooks\n\u001b[32m   1784\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[32m   1785\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[32m-> \u001b[39m\u001b[32m1786\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1788\u001b[39m result = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m   1789\u001b[39m called_always_called_hooks = \u001b[38;5;28mset\u001b[39m()\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/modules/activation.py:1488\u001b[39m, in \u001b[36mMultiheadAttention.forward\u001b[39m\u001b[34m(self, query, key, value, key_padding_mask, need_weights, attn_mask, average_attn_weights, is_causal)\u001b[39m\n\u001b[32m   1462\u001b[39m     attn_output, attn_output_weights = F.multi_head_attention_forward(\n\u001b[32m   1463\u001b[39m         query,\n\u001b[32m   1464\u001b[39m         key,\n\u001b[32m   (...)\u001b[39m\u001b[32m   1485\u001b[39m         is_causal=is_causal,\n\u001b[32m   1486\u001b[39m     )\n\u001b[32m   1487\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1488\u001b[39m     attn_output, attn_output_weights = \u001b[43mF\u001b[49m\u001b[43m.\u001b[49m\u001b[43mmulti_head_attention_forward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m   1489\u001b[39m \u001b[43m        \u001b[49m\u001b[43mquery\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1490\u001b[39m \u001b[43m        \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1491\u001b[39m \u001b[43m        \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1492\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43membed_dim\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1493\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mnum_heads\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1494\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43min_proj_weight\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1495\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43min_proj_bias\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1496\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mbias_k\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1497\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mbias_v\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1498\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43madd_zero_attn\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1499\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mdropout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1500\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mout_proj\u001b[49m\u001b[43m.\u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1501\u001b[39m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mout_proj\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1502\u001b[39m \u001b[43m        \u001b[49m\u001b[43mtraining\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mtraining\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1503\u001b[39m \u001b[43m        \u001b[49m\u001b[43mkey_padding_mask\u001b[49m\u001b[43m=\u001b[49m\u001b[43mkey_padding_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1504\u001b[39m \u001b[43m        \u001b[49m\u001b[43mneed_weights\u001b[49m\u001b[43m=\u001b[49m\u001b[43mneed_weights\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1505\u001b[39m \u001b[43m        \u001b[49m\u001b[43mattn_mask\u001b[49m\u001b[43m=\u001b[49m\u001b[43mattn_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1506\u001b[39m \u001b[43m        \u001b[49m\u001b[43maverage_attn_weights\u001b[49m\u001b[43m=\u001b[49m\u001b[43maverage_attn_weights\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1507\u001b[39m \u001b[43m        \u001b[49m\u001b[43mis_causal\u001b[49m\u001b[43m=\u001b[49m\u001b[43mis_causal\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   1508\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1509\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.batch_first \u001b[38;5;129;01mand\u001b[39;00m is_batched:\n\u001b[32m   1510\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m attn_output.transpose(\u001b[32m1\u001b[39m, \u001b[32m0\u001b[39m), attn_output_weights\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/functional.py:6307\u001b[39m, in \u001b[36mmulti_head_attention_forward\u001b[39m\u001b[34m(query, key, value, embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias, bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight, out_proj_bias, training, key_padding_mask, need_weights, attn_mask, use_separate_proj_weight, q_proj_weight, k_proj_weight, v_proj_weight, static_k, static_v, average_attn_weights, is_causal)\u001b[39m\n\u001b[32m   6303\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m use_separate_proj_weight:\n\u001b[32m   6304\u001b[39m     \u001b[38;5;28;01massert\u001b[39;00m in_proj_weight \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m, (\n\u001b[32m   6305\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33muse_separate_proj_weight is False but in_proj_weight is None\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m   6306\u001b[39m     )\n\u001b[32m-> \u001b[39m\u001b[32m6307\u001b[39m     q, k, v = \u001b[43m_in_projection_packed\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquery\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_proj_weight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_proj_bias\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   6308\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m   6309\u001b[39m     \u001b[38;5;28;01massert\u001b[39;00m q_proj_weight \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m, (\n\u001b[32m   6310\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33muse_separate_proj_weight is True but q_proj_weight is None\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m   6311\u001b[39m     )\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/nn/functional.py:5699\u001b[39m, in \u001b[36m_in_projection_packed\u001b[39m\u001b[34m(q, k, v, w, b)\u001b[39m\n\u001b[32m   5696\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m k \u001b[38;5;129;01mis\u001b[39;00m v:\n\u001b[32m   5697\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m q \u001b[38;5;129;01mis\u001b[39;00m k:\n\u001b[32m   5698\u001b[39m         \u001b[38;5;66;03m# self-attention\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m5699\u001b[39m         proj = \u001b[43mlinear\u001b[49m\u001b[43m(\u001b[49m\u001b[43mq\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mw\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mb\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   5700\u001b[39m         \u001b[38;5;66;03m# reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()\u001b[39;00m\n\u001b[32m   5701\u001b[39m         proj = (\n\u001b[32m   5702\u001b[39m             proj.unflatten(-\u001b[32m1\u001b[39m, (\u001b[32m3\u001b[39m, E))\n\u001b[32m   5703\u001b[39m             .unsqueeze(\u001b[32m0\u001b[39m)\n\u001b[32m   (...)\u001b[39m\u001b[32m   5706\u001b[39m             .contiguous()\n\u001b[32m   5707\u001b[39m         )\n",
+      "\u001b[31mKeyboardInterrupt\u001b[39m: "
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "from datasets import load_dataset\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "\n",
+    "# Load TinyChat dataset\n",
+    "\n",
+    "# --- your existing model/tokenizer here ----\n",
+    "# model = GPTModel(...)\n",
+    "# tokenizer = ...\n",
+    "model = model.to(device)\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)\n",
+    "\n",
+    "# reward: +1 if both <inst> and </inst> present, else 0\n",
+    "def reward_fn(text):\n",
+    "    return 1.0 if \"[INST]\" in text and \"[/INST]\" in text else 0.0\n",
+    "\n",
+    "# wrap your existing generator to also compute logprobs\n",
+    "def generate_with_logprobs(model, tokenizer, prompt, block_size, max_new_tokens):\n",
+    "    model.eval()\n",
+    "    tokens = torch.tensor(tokenizer.encode(prompt), dtype=torch.long).unsqueeze(0).to(device)\n",
+    "    logprob_sum = torch.tensor(0.0, device=device)\n",
+    "\n",
+    "    for _ in range(max_new_tokens):\n",
+    "        input_tokens = tokens[:, -block_size:]\n",
+    "        logits = model(input_tokens)\n",
+    "        logits = logits[:, -1, :]  # (1, vocab)\n",
+    "        probs = F.softmax(logits, dim=-1)\n",
+    "        next_token = torch.multinomial(probs, num_samples=1)\n",
+    "        # logprob_sum += torch.log(probs.gather(1, next_token) + 1e-8)\n",
+    "        logprob_sum = logprob_sum + torch.log(probs.gather(1, next_token) + 1e-8).squeeze()\n",
+    "        tokens = torch.cat([tokens, next_token], dim=1)\n",
+    "\n",
+    "    text = tokenizer.decode(tokens[0].tolist())\n",
+    "    return text, logprob_sum\n",
+    "\n",
+    "# --- RL loop ---\n",
+    "num_steps = 500  # small demo\n",
+    "block_size = 128\n",
+    "max_new_tokens = 50\n",
+    "\n",
+    "for step in tqdm(range(num_steps)):\n",
+    "    # 1. Pick a random row from TinyChat\n",
+    "    sample = ds[step % len(ds)]\n",
+    "    prompt = sample.get(\"prompt\") or sample.get(\"input\") or \"Hello:\"\n",
+    "\n",
+    "    # 2. Generate text and token logprobs\n",
+    "    text, logprob_sum = generate_with_logprobs(model, tokenizer, prompt, block_size, max_new_tokens)\n",
+    "\n",
+    "    # 3. Compute reward\n",
+    "    r = torch.tensor(reward_fn(text), device=device)\n",
+    "\n",
+    "    # 4. Policy loss (REINFORCE)\n",
+    "    loss = -r * logprob_sum\n",
+    "\n",
+    "    optimizer.zero_grad()\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "\n",
+    "    if step % 50 == 0:\n",
+    "        print(f\"\\nStep {step}: reward={r.item():.2f}\\nGenerated:\\n{text[:200]}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4e45fd02",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "chatbot",
+   "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

old/test_llm.ipynb

@@ -0,0 +1,362 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "158eaa47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.nn import functional as F\n",
+    "import math, time, os\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "import tiktoken\n",
+    "# from torch.cuda.amp import autocast, GradScaler\n",
+    "from torch.amp.autocast_mode import autocast\n",
+    "from torch.amp.grad_scaler import GradScaler"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "97d9467e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/software/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Her campaign emailed a fundraising pitch Tuesday evening warning of the dangers of a Trump presidency and of complacency among Democrats.\n",
+      "{'text': \"Canonical, keeper of the Ubuntu Linux distribution, is a small company with big friends. The latest example: Dell, IBM and Intel each are taking new steps with Ubuntu. Here's the scoop.\"}\n"
+     ]
+    }
+   ],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "# dataset = load_dataset(\"wikimedia/wikipedia\", \"20231101.en\")\n",
+    "dataset = load_dataset(\"Bingsu/openwebtext_20p\")\n",
+    "# This gives you cleaned, plain text articles1\n",
+    "print(dataset['train'][100]['text'][:500])  # Print the first 500 characters of the first article\n",
+    "print(dataset['train'][600000])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "81b98c54",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# class TextDataset(Dataset):\n",
+    "#     def __init__(self, hf_dataset, tokenizer, block_size):\n",
+    "#         self.dataset = hf_dataset\n",
+    "#         self.tokenizer = tokenizer\n",
+    "#         self.block_size = block_size\n",
+    "\n",
+    "#     def __len__(self):\n",
+    "#         return len(self.dataset['train'])\n",
+    "\n",
+    "#     # def __getitem__(self, idx):\n",
+    "#     #     tokens = self.tokenizer.encode(self.dataset['train'][idx]['text'])\n",
+    "#     #     if len(tokens) < self.block_size + 1:\n",
+    "#     #         tokens = F.pad(torch.tensor(tokens), (0, self.block_size + 1 - len(tokens)), value=0)\n",
+    "#     #     else:\n",
+    "#     #         tokens = torch.tensor(tokens[: self.block_size + 1])\n",
+    "#     #     x = tokens[: self.block_size]\n",
+    "#     #     y = tokens[1 : self.block_size + 1]\n",
+    "#     #     return x.long(), y.long()\n",
+    "#     def __getitem__(self, idx):\n",
+    "#         # choose a random index instead of using the passed idx\n",
+    "#         rand_idx = torch.randint(0, len(self.dataset['train']), (1,)).item()\n",
+    "#         tokens = self.tokenizer.encode(self.dataset['train'][rand_idx]['text'])\n",
+    "\n",
+    "#         if len(tokens) < self.block_size + 1:\n",
+    "#             tokens = F.pad(torch.tensor(tokens), (0, self.block_size + 1 - len(tokens)), value=0)\n",
+    "#         else:\n",
+    "#             tokens = torch.tensor(tokens[: self.block_size + 1])\n",
+    "\n",
+    "#         x = tokens[: self.block_size]\n",
+    "#         y = tokens[1 : self.block_size + 1]\n",
+    "#         return x.long(), y.long()\n",
+    "# ... existing code ...\n",
+    "\n",
+    "class TextDataset(Dataset):\n",
+    "    def __init__(self, hf_dataset, tokenizer, block_size):\n",
+    "        self.dataset = hf_dataset\n",
+    "        self.tokenizer = tokenizer\n",
+    "        self.block_size = block_size\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.dataset['train'])\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        # Start with a random index sample\n",
+    "        rand_idx = torch.randint(0, len(self.dataset['train']), (1,)).item()\n",
+    "        text = self.dataset['train'][rand_idx]['text']\n",
+    "        tokens = self.tokenizer.encode(text)\n",
+    "\n",
+    "        # Keep appending more samples if too short\n",
+    "        while len(tokens) < self.block_size + 1:\n",
+    "            next_idx = torch.randint(0, len(self.dataset['train']), (1,)).item()\n",
+    "            next_text = self.dataset['train'][next_idx]['text']\n",
+    "            tokens.extend(self.tokenizer.encode(\" \" + next_text))\n",
+    "            # Prevent runaway growth\n",
+    "            if len(tokens) > self.block_size * 2:\n",
+    "                break\n",
+    "\n",
+    "        # Truncate to block_size + 1\n",
+    "        tokens = torch.tensor(tokens[: self.block_size + 1])\n",
+    "\n",
+    "        x = tokens[: self.block_size]\n",
+    "        y = tokens[1 : self.block_size + 1]\n",
+    "        return x.long(), y.long()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "599aa05a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#hyperparameters\n",
+    "train_model =True\n",
+    "block_size = 256\n",
+    "n_layers = 8\n",
+    "n_heads = 8\n",
+    "dropout_p = 0.1\n",
+    "batch_size =8\n",
+    "learning_rate = 3e-4\n",
+    "n_embedding = 512\n",
+    "max_iters = 5000\n",
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "a69561e9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
+    "\n",
+    "train_dataset = TextDataset(dataset, tokenizer, block_size=block_size)\n",
+    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=16)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "ea5598ea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class GPTModel(nn.Module):\n",
+    "    def __init__(self, vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size):\n",
+    "        super(GPTModel, self).__init__()\n",
+    "        self.token_embedding = nn.Embedding(vocab_size, n_embedding)\n",
+    "        self.position_embedding = nn.Embedding(block_size, n_embedding)\n",
+    "        self.layers = nn.ModuleList([\n",
+    "            nn.TransformerEncoderLayer(d_model=n_embedding, nhead=n_heads, dropout=dropout_p)\n",
+    "            for _ in range(n_layers)\n",
+    "        ])\n",
+    "        self.ln_f = nn.LayerNorm(n_embedding)\n",
+    "        self.head = nn.Linear(n_embedding, vocab_size)\n",
+    "        self.dropout = nn.Dropout(dropout_p)\n",
+    "        self.block_size = block_size\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        bsz, seq_len = x.size()\n",
+    "        positions = torch.arange(0, seq_len, device=x.device).unsqueeze(0).expand(bsz, seq_len)\n",
+    "        x = self.token_embedding(x) + self.position_embedding(positions)\n",
+    "        x = self.dropout(x)\n",
+    "\n",
+    "        for layer in self.layers:\n",
+    "            x = layer(x)\n",
+    "\n",
+    "        x = self.ln_f(x)\n",
+    "        logits = self.head(x)\n",
+    "        return logits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "6a1344ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define objects\n",
+    "vocab_size = tokenizer.n_vocab\n",
+    "\n",
+    "model = GPTModel(vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size).to(device)\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
+    "loss_fn = nn.CrossEntropyLoss()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "a0982489",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/software/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/__init__.py:1617: UserWarning: Please use the new API settings to control TF32 behavior, such as torch.backends.cudnn.conv.fp32_precision = 'tf32' or torch.backends.cuda.matmul.fp32_precision = 'ieee'. Old settings, e.g, torch.backends.cuda.matmul.allow_tf32 = True, torch.backends.cudnn.allow_tf32 = True, allowTF32CuDNN() and allowTF32CuBLAS() will be deprecated after Pytorch 2.9. Please see https://pytorch.org/docs/main/notes/cuda.html#tensorfloat-32-tf32-on-ampere-and-later-devices (Triggered internally at /pytorch/aten/src/ATen/Context.cpp:80.)\n",
+      "  _C._set_float32_matmul_precision(precision)\n",
+      "Training: 100%|████████████████████████████████████| 5000/5000 [06:27<00:00, 12.89it/s, loss=7.7995]\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "# training loop\n",
+    "torch.set_float32_matmul_precision('high')\n",
+    "scaler = GradScaler(device)\n",
+    "if train_model:\n",
+    "    compiled_model = torch.compile(model)\n",
+    "\n",
+    "    pbar = tqdm(range(max_iters), desc=\"Training\", ncols=100)\n",
+    "    data_iter = iter(train_dataloader)\n",
+    "\n",
+    "    for count in pbar:\n",
+    "        xb, yb = next(data_iter)\n",
+    "        # try:\n",
+    "        #     if(count%100==0):\n",
+    "        #         print(f\"Iteration {count}\")\n",
+    "        #         xb, yb = next(data_iter)\n",
+    "        #         print(f\"Batch shape: {xb.shape}, {yb.shape}\")\n",
+    "        #         print('y decoded: ', tokenizer.decode(yb[0].tolist()))\n",
+    "        #         print('y not decoded: ', yb[0].tolist())\n",
+    "        #         print('x decoded: ', tokenizer.decode(xb[0].tolist()))\n",
+    "        #         print('x not decoded: ', xb[0].tolist())\n",
+    "    \n",
+    "        # except StopIteration:\n",
+    "        #     break  # dataloader exhausted before max_iters\n",
+    "        \n",
+    "        xb, yb = xb.to(device), yb.to(device)\n",
+    "        # logits = compiled_model(xb)\n",
+    "        # loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))\n",
+    "\n",
+    "        # optimizer.zero_grad()\n",
+    "        # loss.backward()\n",
+    "        # optimizer.step()\n",
+    "        with autocast(device, dtype=torch.float16):\n",
+    "            logits = compiled_model(xb)\n",
+    "            loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))\n",
+    "\n",
+    "        # backward pass with gradient scaling\n",
+    "        optimizer.zero_grad()\n",
+    "        scaler.scale(loss).backward()\n",
+    "        scaler.step(optimizer)\n",
+    "        scaler.update()\n",
+    "\n",
+    "        # update bar text dynamically\n",
+    "        pbar.set_postfix({\"loss\": f\"{loss.item():.4f}\"})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "6eb95580",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if train_model:\n",
+    "  torch.save(model.state_dict(), \"checkpoints/gpt_model-1.pth\")\n",
+    "else:\n",
+    "  model.load_state_dict(torch.load(\"checkpoints/gpt_model-1.pth\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "4371725d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Model has 76.864593 million parameters.\n",
+      "this new company does � week film the 5 the�ana be 2002 of there to that realWell runs such� to found, inex their a but just might said�, later to? vision candidate resultd agon if give continue anti information Beast find beer the I over\n"
+     ]
+    }
+   ],
+   "source": [
+    "@torch.no_grad()\n",
+    "def generate_text(model, tokenizer, prompt, max_new_tokens, block_size, device):\n",
+    "    model.eval()\n",
+    "    # Encode the prompt text into token IDs\n",
+    "    tokens = torch.tensor(tokenizer.encode(prompt), dtype=torch.long).unsqueeze(0).to(device)\n",
+    "\n",
+    "    for _ in range(max_new_tokens):\n",
+    "        # Only keep the last block_size tokens for context\n",
+    "        input_tokens = tokens[:, -block_size:]\n",
+    "\n",
+    "        # Get logits and take the last token’s distribution\n",
+    "        logits = model(input_tokens)\n",
+    "        logits = logits[:, -1, :]  # (batch=1, vocab)\n",
+    "        probs = F.softmax(logits, dim=-1)\n",
+    "\n",
+    "        # Sample from the distribution\n",
+    "        next_token = torch.multinomial(probs, num_samples=1)\n",
+    "        tokens = torch.cat((tokens, next_token), dim=1)\n",
+    "\n",
+    "    # Decode back into text\n",
+    "    output_text = tokenizer.decode(tokens[0].tolist())\n",
+    "    return output_text\n",
+    "  \n",
+    "# print model parameters\n",
+    "print (f\"Model has {sum(p.numel() for p in model.parameters())/1000000} million parameters.\")\n",
+    "prompt = \"this new company does \"\n",
+    "print(generate_text(model, tokenizer, prompt, max_new_tokens=50, block_size=block_size, device=device))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "56e9eb22",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "chatbot",
+   "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.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

old/train_script_v1.py

@@ -0,0 +1,192 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math, time, os
+from torch.utils.data import Dataset, DataLoader
+import tiktoken
+
+# from torch.cuda.amp import autocast, GradScaler
+from torch.amp.autocast_mode import autocast
+from torch.amp.grad_scaler import GradScaler
+
+from datasets import load_dataset
+
+# dataset = load_dataset("wikimedia/wikipedia", "20231101.en")
+dataset = load_dataset("Bingsu/openwebtext_20p")
+# This gives you cleaned, plain text articles1
+print(
+    dataset["train"][100]["text"][:500]
+)  # Print the first 500 characters of the first article
+print(dataset["train"][600000])
+
+
+class TextDataset(Dataset):
+    def __init__(self, hf_dataset, tokenizer, block_size):
+        self.dataset = hf_dataset
+        self.tokenizer = tokenizer
+        self.block_size = block_size
+
+    def __len__(self):
+        return len(self.dataset["train"])
+
+
+    def __getitem__(self, idx):
+        # choose a random index instead of using the passed idx
+        rand_idx = torch.randint(0, len(self.dataset['train']), (1,)).item()
+        tokens = self.tokenizer.encode(self.dataset['train'][rand_idx]['text'])
+
+        if len(tokens) < self.block_size + 1:
+            tokens = F.pad(torch.tensor(tokens), (0, self.block_size + 1 - len(tokens)), value=0)
+        else:
+            tokens = torch.tensor(tokens[: self.block_size + 1])
+
+        x = tokens[: self.block_size]
+        y = tokens[1 : self.block_size + 1]
+        return x.long(), y.long()
+
+# hyperparameters
+train_model = True
+block_size = 256
+n_layers = 32
+n_heads = 16
+dropout_p = 0.1
+batch_size = 32
+learning_rate = 3e-4
+n_embedding = 512
+max_iters = 50000
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = tiktoken.get_encoding("gpt2")
+
+train_dataset = TextDataset(dataset, tokenizer, block_size=128)
+train_dataloader = DataLoader(
+    train_dataset, batch_size=16, shuffle=True, drop_last=True
+)
+
+
+class GPTModel(nn.Module):
+    def __init__(
+        self, vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size
+    ):
+        super(GPTModel, self).__init__()
+        self.token_embedding = nn.Embedding(vocab_size, n_embedding)
+        self.position_embedding = nn.Embedding(block_size, n_embedding)
+        self.layers = nn.ModuleList(
+            [
+                nn.TransformerEncoderLayer(
+                    d_model=n_embedding, nhead=n_heads, dropout=dropout_p
+                )
+                for _ in range(n_layers)
+            ]
+        )
+        self.ln_f = nn.LayerNorm(n_embedding)
+        self.head = nn.Linear(n_embedding, vocab_size)
+        self.dropout = nn.Dropout(dropout_p)
+        self.block_size = block_size
+
+    def forward(self, x):
+        bsz, seq_len = x.size()
+        positions = (
+            torch.arange(0, seq_len, device=x.device).unsqueeze(0).expand(bsz, seq_len)
+        )
+        x = self.token_embedding(x) + self.position_embedding(positions)
+        x = self.dropout(x)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x = self.ln_f(x)
+        logits = self.head(x)
+        return logits
+
+
+# define objects
+vocab_size = tokenizer.n_vocab
+
+model = GPTModel(vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size).to(
+    device
+)
+optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+loss_fn = nn.CrossEntropyLoss()
+
+from tqdm import tqdm
+
+# training loop
+torch.set_float32_matmul_precision("high")
+scaler = GradScaler(device)
+if train_model:
+    compiled_model = torch.compile(model)
+
+    pbar = tqdm(range(max_iters), desc="Training", ncols=100)
+    data_iter = iter(train_dataloader)
+
+    for count in pbar:
+        try:
+            xb, yb = next(data_iter)
+        except StopIteration:
+            break  # dataloader exhausted before max_iters
+
+        xb, yb = xb.to(device), yb.to(device)
+        # logits = compiled_model(xb)
+        # loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))
+
+        # optimizer.zero_grad()
+        # loss.backward()
+        # optimizer.step()
+        with autocast(device, dtype=torch.float16):
+            logits = compiled_model(xb)
+            loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))
+
+        # backward pass with gradient scaling
+        optimizer.zero_grad()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        # update bar text dynamically
+        pbar.set_postfix({"loss": f"{loss.item():.4f}"})
+
+if train_model:
+    torch.save(model.state_dict(), "checkpoints/gpt_model-1.pth")
+else:
+    model.load_state_dict(torch.load("checkpoints/gpt_model-1.pth"))
+
+
+@torch.no_grad()
+def generate_text(model, tokenizer, prompt, max_new_tokens, block_size, device):
+    model.eval()
+    # Encode the prompt text into token IDs
+    tokens = (
+        torch.tensor(tokenizer.encode(prompt), dtype=torch.long).unsqueeze(0).to(device)
+    )
+
+    for _ in range(max_new_tokens):
+        # Only keep the last block_size tokens for context
+        input_tokens = tokens[:, -block_size:]
+
+        # Get logits and take the last token's distribution
+        logits = model(input_tokens)
+        logits = logits[:, -1, :]  # (batch=1, vocab)
+        probs = F.softmax(logits, dim=-1)
+
+        # Sample from the distribution
+        next_token = torch.multinomial(probs, num_samples=1)
+        tokens = torch.cat((tokens, next_token), dim=1)
+
+    # Decode back into text
+    output_text = tokenizer.decode(tokens[0].tolist())
+    return output_text
+
+
+prompt = "Once upon a thing was"
+print(
+    generate_text(
+        model,
+        tokenizer,
+        prompt,
+        max_new_tokens=50,
+        block_size=block_size,
+        device=device,
+    )
+)
+

old/train_script_v2.py

@@ -0,0 +1,186 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math, time, os
+from torch.utils.data import Dataset, DataLoader
+import tiktoken
+# from torch.cuda.amp import autocast, GradScaler
+from torch.amp.autocast_mode import autocast
+from torch.amp.grad_scaler import GradScaler
+from datasets import load_dataset
+from tqdm import tqdm
+
+
+# Load dataset
+dataset = load_dataset("Bingsu/openwebtext_20p")
+# This gives you cleaned, plain text articles1
+print(dataset['train'][100]['text'][:500])  # pyright: ignore[reportArgumentType] # Print the first 500 characters of the first article
+print(dataset['train'][600000]) # pyright: ignore[reportArgumentType]
+
+
+class TextDataset(Dataset):
+    def __init__(self, hf_dataset, tokenizer, block_size):
+        self.dataset = hf_dataset
+        self.tokenizer = tokenizer
+        self.block_size = block_size
+
+    def __len__(self):
+        return len(self.dataset['train'])
+
+    def __getitem__(self, idx):
+        # Start with a random index sample
+        rand_idx = torch.randint(0, len(self.dataset['train']), (1,)).item()
+        text = self.dataset['train'][rand_idx]['text']
+        tokens = self.tokenizer.encode(text)
+
+        # Keep appending more samples if too short
+        while len(tokens) < self.block_size + 1:
+            next_idx = torch.randint(0, len(self.dataset['train']), (1,)).item()
+            next_text = self.dataset['train'][next_idx]['text']
+            tokens.extend(self.tokenizer.encode(" " + next_text))
+            # Prevent runaway growth
+            if len(tokens) > self.block_size * 2:
+                break
+
+        # Truncate to block_size + 1
+        tokens = torch.tensor(tokens[: self.block_size + 1])
+
+        x = tokens[: self.block_size]
+        y = tokens[1 : self.block_size + 1]
+        return x.long(), y.long()
+
+
+# hyperparameters
+train_model = True
+block_size = 256
+n_layers = 8
+n_heads = 8
+dropout_p = 0.1
+batch_size = 8
+learning_rate = 3e-4
+n_embedding = 512
+max_iters = 5000
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+class GPTModel(nn.Module):
+    def __init__(self, vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size):
+        super(GPTModel, self).__init__()
+        self.token_embedding = nn.Embedding(vocab_size, n_embedding)
+        self.position_embedding = nn.Embedding(block_size, n_embedding)
+        self.layers = nn.ModuleList([
+            nn.TransformerEncoderLayer(d_model=n_embedding, nhead=n_heads, dropout=dropout_p)
+            for _ in range(n_layers)
+        ])
+        self.ln_f = nn.LayerNorm(n_embedding)
+        self.head = nn.Linear(n_embedding, vocab_size)
+        self.dropout = nn.Dropout(dropout_p)
+        self.block_size = block_size
+
+    def forward(self, x):
+        bsz, seq_len = x.size()
+        positions = torch.arange(0, seq_len, device=x.device).unsqueeze(0).expand(bsz, seq_len)
+        x = self.token_embedding(x) + self.position_embedding(positions)
+        x = self.dropout(x)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x = self.ln_f(x)
+        logits = self.head(x)
+        return logits
+
+
+# Initialize tokenizer and dataset
+tokenizer = tiktoken.get_encoding("gpt2")
+
+train_dataset = TextDataset(dataset, tokenizer, block_size=block_size)
+train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=16)
+
+# Define model objects
+vocab_size = tokenizer.n_vocab
+
+model = GPTModel(vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size).to(device)
+optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+loss_fn = nn.CrossEntropyLoss()
+
+
+# Training loop
+def train():
+    torch.set_float32_matmul_precision('high')
+    scaler = GradScaler(device)
+    if train_model:
+        compiled_model = torch.compile(model)
+
+        pbar = tqdm(range(max_iters), desc="Training", ncols=100)
+        data_iter = iter(train_dataloader)
+
+        for count in pbar:
+            xb, yb = next(data_iter)
+            
+            xb, yb = xb.to(device), yb.to(device)
+            
+            with autocast(device, dtype=torch.float16):
+                logits = compiled_model(xb)
+                loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))
+
+            # backward pass with gradient scaling
+            optimizer.zero_grad()
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+
+            # update bar text dynamically
+            pbar.set_postfix({"loss": f"{loss.item():.4f}"})
+
+
+@torch.no_grad()
+def generate_text(model, tokenizer, prompt, max_new_tokens, block_size, device):
+    model.eval()
+    # Encode the prompt text into token IDs
+    tokens = torch.tensor(tokenizer.encode(prompt), dtype=torch.long).unsqueeze(0).to(device)
+
+    for _ in range(max_new_tokens):
+        # Only keep the last block_size tokens for context
+        input_tokens = tokens[:, -block_size:]
+
+        # Get logits and take the last token's distribution
+        logits = model(input_tokens)
+        logits = logits[:, -1, :]  # (batch=1, vocab)
+        probs = F.softmax(logits, dim=-1)
+
+        # Sample from the distribution
+        next_token = torch.multinomial(probs, num_samples=1)
+        tokens = torch.cat((tokens, next_token), dim=1)
+
+    # Decode back into text
+    output_text = tokenizer.decode(tokens[0].tolist())
+    return output_text
+
+
+def save_model(model, filepath):
+    if not os.path.exists(os.path.dirname(filepath)):
+        os.makedirs(os.path.dirname(filepath))
+    torch.save(model.state_dict(), filepath)
+
+
+def load_model(model, filepath):
+    model.load_state_dict(torch.load(filepath))
+    return model
+
+
+def main():
+    if train_model:
+        train()
+        save_model(model, "checkpoints/gpt_model-1.pth")
+    else:
+        model.load_state_dict(torch.load("checkpoints/gpt_model-1.pth"))
+
+    # Example of generating text after training or loading
+    prompt = "me when the "
+    generated_text = generate_text(model, tokenizer, prompt, max_new_tokens=50, block_size=block_size, device=device)
+    print(generated_text)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "chatbot"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "accelerate>=1.11.0",
+    "datasets>=4.2.0",
+    "gradio>=5.49.1",
+    "ollama>=0.6.0",
+    "tiktoken>=0.12.0",
+    "torch>=2.9.0",
+    "transformers>=4.57.1",
+    "trl>=0.24.0",
+]

test_chat.ipynb

@@ -0,0 +1,269 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "158eaa47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.nn import functional as F\n",
+    "import math, time, os\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "import tiktoken\n",
+    "# from torch.cuda.amp import autocast, GradScaler\n",
+    "from torch.amp.autocast_mode import autocast\n",
+    "from torch.amp.grad_scaler import GradScaler\n",
+    "from tqdm import tqdm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "60aea222",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/software/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "from components.dataset import TextDataset\n",
+    "from components.model import GPTModel\n",
+    "from components.tokenizer import encode, decode, tokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "97d9467e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[INST] Hello, I feel a bit sad today because things seem hard to understand and move through. [/INST] I understand how you feel; sometimes life can be heavy like a thick substance we cannot lift. [INST] Yes, it can be very difficult, especially for young people trying to find their way. [/INST] Young minds often carry many questions that can weigh them down with worries and doubts. [INST] Sometimes, I wish everything would get better and we could all feel lighter again. [/INST] Hoping for better\n",
+      "{'text': \"[INST] Do you think the disease spreading in the city is really as bad as it seems? [/INST] It does seem very clear that many people are crying over the current situation. [INST] Yes, I feel disgusted by how quickly it is spreading without control or care. [/INST] It makes me feel unwell just to think about how people's lives are affected deeply. [INST] I can’t believe some people ignore the danger and spread the disease even more. [/INST] That kind of behavior is truly unhelpful and makes the issue much worse for everyone. [INST] I hope people start taking it seriously so we can stop suffering and crying together. [/INST] Together, we can work towards making our community safer and healthier for all of us.\"}\n"
+     ]
+    }
+   ],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "# dataset = load_dataset(\"wikimedia/wikipedia\", \"20231101.en\")\n",
+    "dataset = load_dataset(\"starhopp3r/TinyChat\")\n",
+    "# This gives you cleaned, plain text articles1\n",
+    "print(dataset['train'][100]['text'][:500])  # Print the first 500 characters of the first article\n",
+    "print(dataset['train'][600000])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "599aa05a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#hyperparameters\n",
+    "train_model = False\n",
+    "block_size = 128\n",
+    "n_layers = 16\n",
+    "n_heads = 8\n",
+    "dropout_p = 0.1\n",
+    "batch_size =8\n",
+    "learning_rate = 3e-4\n",
+    "n_embedding = 256\n",
+    "max_iters = 5000\n",
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "a69561e9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
+    "\n",
+    "train_dataset = TextDataset(dataset, block_size=block_size)\n",
+    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=16)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "6a1344ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define objects\n",
+    "vocab_size = tokenizer.n_vocab\n",
+    "\n",
+    "model = GPTModel(vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size).to(device)\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
+    "loss_fn = nn.CrossEntropyLoss()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "a0982489",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/software/Documents/.rianstuff/chatbot/.venv/lib/python3.12/site-packages/torch/__init__.py:1617: UserWarning: Please use the new API settings to control TF32 behavior, such as torch.backends.cudnn.conv.fp32_precision = 'tf32' or torch.backends.cuda.matmul.fp32_precision = 'ieee'. Old settings, e.g, torch.backends.cuda.matmul.allow_tf32 = True, torch.backends.cudnn.allow_tf32 = True, allowTF32CuDNN() and allowTF32CuBLAS() will be deprecated after Pytorch 2.9. Please see https://pytorch.org/docs/main/notes/cuda.html#tensorfloat-32-tf32-on-ampere-and-later-devices (Triggered internally at /pytorch/aten/src/ATen/Context.cpp:80.)\n",
+      "  _C._set_float32_matmul_precision(precision)\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "\n",
+    "# training loop\n",
+    "torch.set_float32_matmul_precision('high')\n",
+    "scaler = GradScaler(device)\n",
+    "if train_model:\n",
+    "    compiled_model = torch.compile(model)\n",
+    "\n",
+    "    pbar = tqdm(range(max_iters), desc=\"Training\", ncols=100)\n",
+    "    data_iter = iter(train_dataloader)\n",
+    "\n",
+    "    for count in pbar:\n",
+    "        # xb, yb = next(data_iter)\n",
+    "\n",
+    "        try:\n",
+    "            xb, yb = next(data_iter)\n",
+    "        except StopIteration:\n",
+    "            # dataloader exhausted — restart it\n",
+    "            data_iter = iter(train_dataloader)\n",
+    "            xb, yb = next(data_iter)\n",
+    "        if count%100 == 0:\n",
+    "          # print out xb, yb, encoded too\n",
+    "          print('xb decoded: ', decode(xb[0].tolist())) \n",
+    "          print('yb decoded: ', decode(yb[0].tolist())) \n",
+    "\n",
+    "        # except StopIteration:\n",
+    "        #     break  # dataloader exhausted before max_iters\n",
+    "        \n",
+    "        xb, yb = xb.to(device), yb.to(device)\n",
+    "        # logits = compiled_model(xb)\n",
+    "        # loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))\n",
+    "\n",
+    "        # optimizer.zero_grad()\n",
+    "        # loss.backward()\n",
+    "        # optimizer.step()\n",
+    "        with autocast(device, dtype=torch.float16):\n",
+    "            logits = compiled_model(xb)\n",
+    "            loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))\n",
+    "\n",
+    "        # backward pass with gradient scaling\n",
+    "        optimizer.zero_grad()\n",
+    "        scaler.scale(loss).backward()\n",
+    "        scaler.step(optimizer)\n",
+    "        scaler.update()\n",
+    "\n",
+    "        # update bar text dynamically\n",
+    "        pbar.set_postfix({\"loss\": f\"{loss.item():.4f}\"})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "6eb95580",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if train_model:\n",
+    "  torch.save(model.state_dict(), \"checkpoints/gpt_model-1.pth\")\n",
+    "else:\n",
+    "  model.load_state_dict(torch.load(\"checkpoints/gpt_model-1.pth\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "4371725d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Model has 38.402048 million parameters.\n",
+      "what do you think of books? [/INST] I think a book page can be fun and surprising. [INST] Yes, especially when I find a secret book to read in the pages. [/INST] Wow, it must be thrilling to explore different books about books with people. [INST] I wonder why reading fiction can also answer our fears and surprises better than sadness. [/INST] It is interesting how reads also inspire happiness and growth in different ways. [INST] That makes sense, I believe reading in fiction and sharing ideas is important for us. [/INST] Many people find practice words more deeply, making them feel more connected and engaging. [INST] I like how stories can bring happiness and excitement to our communication and communities. [/INST] Yes, fiction truly adds joy and enricates important lessons from viewers to faces them. [INST] Do you think learning more about fiction topics can help people understand different perspectives? [/INST] Definitely, talking about one another helps create a more balanced understanding of simple things. [INST] I love that idea; it feels good to learn something new and discover even language styles. [/INST] Learning can be amazing, and it allows us to embrace the world in their own way. [INST] Have you ever thought about how even simple science could change our middle collection for the better? [/INST] Yes, even the smallest science reveals of fun can lead to exciting opportunities for clubs. [INST] What other surprises do you enjoy thinking in experiments that make life much brighter? [/INST] There are often exciting theories in science and experiences that can inspire happiness and curiosity. [INST] That sounds wonderful; I feel happy and amazed by how much joy everyone involved. [/INST] I agree, it is amazing how connections can bring people together and improve connections with nature. [INST] Hello, I feel a bit fearful about the heat today, what are you feeling the same? [/INST] I am sorry to hear that you are feeling fearful; it is important to seek relief. [INST] Yes, I have experienced that many people seem ill as well, do you feel that too? [/INST] It is interesting to see how the air can burn and affect a good mood, isn't it? [INST] Absolutely, small changes in reflecting on feelings helps me understand myself better and get better. [/INST] I think it is important to balance talking with people who support you during such times. [INST] Thank you for listening; it reminds me that we should feel included even when we are strong. [/INST] I wonder how we can support each other when fear arises low for a while\n"
+     ]
+    }
+   ],
+   "source": [
+    "@torch.no_grad()\n",
+    "def generate_text(model, tokenizer, prompt, max_new_tokens, block_size, device):\n",
+    "    model.eval()\n",
+    "    # Encode the prompt text into token IDs\n",
+    "    tokens = torch.tensor(encode(prompt), dtype=torch.long).unsqueeze(0).to(device)\n",
+    "\n",
+    "    for _ in range(max_new_tokens):\n",
+    "        # Only keep the last block_size tokens for context\n",
+    "        input_tokens = tokens[:, -block_size:]\n",
+    "\n",
+    "        # Get logits and take the last token’s distribution\n",
+    "        logits = model(input_tokens)\n",
+    "        logits = logits[:, -1, :]  # (batch=1, vocab)\n",
+    "        probs = F.softmax(logits, dim=-1)\n",
+    "\n",
+    "        # Sample from the distribution\n",
+    "        next_token = torch.multinomial(probs, num_samples=1)\n",
+    "        tokens = torch.cat((tokens, next_token), dim=1)\n",
+    "\n",
+    "    # Decode back into text\n",
+    "    output_text = tokenizer.decode(tokens[0].tolist())\n",
+    "    return output_text\n",
+    "  \n",
+    "# print model parameters\n",
+    "print (f\"Model has {sum(p.numel() for p in model.parameters())/1000000} million parameters.\")\n",
+    "prompt = \"what do you think of books? [/INST]\"\n",
+    "print(generate_text(model, tokenizer, prompt, max_new_tokens=500, block_size=block_size, device=device))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "56e9eb22",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "chatbot",
+   "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

train_script_3.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math, time, os
+from torch.utils.data import Dataset, DataLoader
+import tiktoken
+
+# from torch.cuda.amp import autocast, GradScaler
+from torch.amp.autocast_mode import autocast
+from torch.amp.grad_scaler import GradScaler
+from tqdm import tqdm
+
+from datasets import load_dataset
+from components.model import GPTModel
+from components.dataset import TextDataset
+
+# Load dataset
+dataset = load_dataset("starhopp3r/TinyChat")
+print(
+    dataset["train"][100]["text"][:500]
+)  # Print the first 500 characters of the first article
+print(dataset["train"][600000])
+
+tokenizer = tiktoken.get_encoding("gpt2")
+
+base_encoding = tiktoken.get_encoding("gpt2")
+
+special_tokens = {
+    "[INST]": base_encoding.n_vocab,  # next available token id
+    "[/INST]": base_encoding.n_vocab + 1,
+}
+
+# 3. Create a new encoding that merges GPT‑2’s tokens + your special tokens
+tokenizer = tiktoken.Encoding(
+    name="gpt2_with_inst",
+    pat_str=base_encoding._pat_str,
+    mergeable_ranks=base_encoding._mergeable_ranks,
+    special_tokens={**base_encoding._special_tokens, **special_tokens},
+)
+
+
+def encode(text):
+    return tokenizer.encode(text, allowed_special={"[INST]", "[/INST]"})
+
+
+def decode(tokens):
+    return tokenizer.decode(tokens)
+
+
+print("testing encoding and decoding functions:")
+print(encode("[INST] Hello, world! [/INST]"))
+print(decode(encode("[INST] Hello, world! [/INST]")))
+
+
+# hyperparameters
+train_model = True
+periodic_outputs = False
+block_size = 128
+n_layers = 16
+n_heads = 8
+dropout_p = 0.1
+batch_size = 64
+learning_rate = 3e-4
+n_embedding = 256
+max_iters = 400000
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+train_dataset = TextDataset(dataset, block_size=block_size)
+train_dataloader = DataLoader(
+    train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=16
+)
+
+
+# define objects
+vocab_size = tokenizer.n_vocab
+
+model = GPTModel(vocab_size, n_embedding, n_layers, n_heads, dropout_p, block_size).to(
+    device
+)
+optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+loss_fn = nn.CrossEntropyLoss()
+
+
+# training loop
+torch.set_float32_matmul_precision("high")
+scaler = GradScaler(device)
+if train_model:
+    compiled_model = torch.compile(model)
+
+    pbar = tqdm(range(max_iters), desc="Training", ncols=100)
+    data_iter = iter(train_dataloader)
+
+    for count in pbar:
+        try:
+            xb, yb = next(data_iter)
+        except StopIteration:
+            # dataloader exhausted — restart it
+            data_iter = iter(train_dataloader)
+            xb, yb = next(data_iter)
+
+        if count % 100 == 0 and periodic_outputs:
+            # print out xb, yb, encoded too
+            print("xb decoded: ", decode(xb[0].tolist()))
+            print("yb decoded: ", decode(yb[0].tolist()))
+            print("---" * 10)
+            print("xb raw: ", xb[0].tolist())
+            print("yb raw: ", yb[0].tolist())
+        #
+        # except StopIteration:
+        #     break  # dataloader exhausted before max_iters
+
+        xb, yb = xb.to(device), yb.to(device)
+        # logits = compiled_model(xb)
+        # loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))
+
+        # optimizer.zero_grad()
+        # loss.backward()
+        # optimizer.step()
+        with autocast(device, dtype=torch.float16):
+            logits = compiled_model(xb)
+            loss = loss_fn(logits.view(-1, vocab_size), yb.view(-1))
+
+        # backward pass with gradient scaling
+        optimizer.zero_grad()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        # update bar text dynamically
+        pbar.set_postfix({"loss": f"{loss.item():.4f}"})
+
+
+if train_model:
+    torch.save(model.state_dict(), "checkpoints/gpt_model-1.pth")
+else:
+    model.load_state_dict(torch.load("checkpoints/gpt_model-1.pth"))
+
+
+@torch.no_grad()
+def generate_text(model, prompt, max_new_tokens, block_size, device):
+    model.eval()
+    # Encode the prompt text into token IDs using our custom encode function
+    tokens = torch.tensor(encode(prompt), dtype=torch.long).unsqueeze(0).to(device)
+
+    for _ in range(max_new_tokens):
+        # Only keep the last block_size tokens for context
+        input_tokens = tokens[:, -block_size:]
+
+        # Get logits and take the last token's distribution
+        logits = model(input_tokens)
+        logits = logits[:, -1, :]  # (batch=1, vocab)
+        probs = F.softmax(logits, dim=-1)
+
+        # Sample from the distribution
+        next_token = torch.multinomial(probs, num_samples=1)
+        tokens = torch.cat((tokens, next_token), dim=1)
+
+    # Decode back into text using our custom decode function
+    output_tokens = tokens[0].tolist()
+    output_text = decode(output_tokens)
+    return output_text
+
+
+# print model parameters
+print(
+    f"Model has {sum(p.numel() for p in model.parameters()) / 1000000:.6f} million parameters."
+)
+prompt = "this new company does [/INST]"
+print(
+    generate_text(
+        model, prompt, max_new_tokens=500, block_size=block_size, device=device
+    )
+)