Upload main.ipynb with huggingface_hub

main.ipynb

@@ -2161,6 +2161,482 @@
     "\n",
     "visualize_diffusion(model, tokenizer)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ft_header",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "---\n",
+    "# Part 2: Fine-tuning for Chat\n",
+    "\n",
+    "Now we turn the pretrained MDLM into a **chatbot** using supervised fine-tuning on dialogue data.\n",
+    "\n",
+    "## How diffusion chat works\n",
+    "1. Format: `<|user|> message <|assistant|> response <|end|>`\n",
+    "2. **Training**: Mask only the response tokens \u2014 the user message stays visible as context\n",
+    "3. **Inference**: User types a message \u2192 freeze those tokens \u2192 diffusion unmasks only the response\n",
+    "4. **The cool part**: The response materializes all at once, not left-to-right"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "ft_config",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ============================================================\n",
+    "# FINE-TUNING CONFIG\n",
+    "# ============================================================\n",
+    "\n",
+    "@dataclass\n",
+    "class FinetuneConfig:\n",
+    "    # Training\n",
+    "    ft_steps: int = 5000\n",
+    "    ft_batch_size: int = 16\n",
+    "    ft_lr: float = 5e-5          # Lower LR for fine-tuning\n",
+    "    ft_warmup: int = 200\n",
+    "    max_response_len: int = 128   # Max response length\n",
+    "    max_prompt_len: int = 64      # Max prompt length\n",
+    "    log_every: int = 50\n",
+    "    sample_every: int = 500\n",
+    "\n",
+    "ft_config = FinetuneConfig()\n",
+    "\n",
+    "# Add special tokens to tokenizer\n",
+    "SPECIAL_TOKENS = {\n",
+    "    'additional_special_tokens': ['<|user|>', '<|assistant|>', '<|end|>']\n",
+    "}\n",
+    "tokenizer.add_special_tokens(SPECIAL_TOKENS)\n",
+    "\n",
+    "USER_TOKEN = tokenizer.convert_tokens_to_ids('<|user|>')\n",
+    "ASST_TOKEN = tokenizer.convert_tokens_to_ids('<|assistant|>')\n",
+    "END_TOKEN = tokenizer.convert_tokens_to_ids('<|end|>')\n",
+    "\n",
+    "print(f'Special token IDs: USER={USER_TOKEN}, ASST={ASST_TOKEN}, END={END_TOKEN}')\n",
+    "\n",
+    "# Resize model embeddings to accommodate new tokens\n",
+    "old_vocab = config.vocab_size\n",
+    "new_vocab = len(tokenizer)\n",
+    "if new_vocab > old_vocab:\n",
+    "    # Expand embedding and output projection\n",
+    "    old_emb = model_unwrapped.token_emb.weight.data\n",
+    "    model_unwrapped.token_emb = nn.Embedding(new_vocab, config.hidden_dim).to(device)\n",
+    "    model_unwrapped.token_emb.weight.data[:old_vocab] = old_emb\n",
+    "    # Re-tie output projection\n",
+    "    model_unwrapped.output_proj = nn.Linear(config.hidden_dim, new_vocab, bias=False).to(device)\n",
+    "    model_unwrapped.output_proj.weight = model_unwrapped.token_emb.weight\n",
+    "    # Update config\n",
+    "    config.vocab_size = new_vocab\n",
+    "    model_unwrapped.config.vocab_size = new_vocab\n",
+    "    print(f'Resized embeddings: {old_vocab} -> {new_vocab}')\n",
+    "\n",
+    "print(f'Fine-tune config ready')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "ft_dataset",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ============================================================\n",
+    "# DIALOGUE DATASET\n",
+    "# ============================================================\n",
+    "\n",
+    "from datasets import load_dataset\n",
+    "\n",
+    "# Using Alpaca-cleaned: simple instruction-response pairs\n",
+    "print('Loading Alpaca dataset...')\n",
+    "alpaca = load_dataset('yahma/alpaca-cleaned', split='train')\n",
+    "print(f'Loaded {len(alpaca)} examples')\n",
+    "\n",
+    "class ChatDataset(torch.utils.data.Dataset):\n",
+    "    \"\"\"Format dialogue as: <|user|> instruction <|assistant|> response <|end|>\n",
+    "    \n",
+    "    Returns:\n",
+    "        input_ids: full sequence token ids\n",
+    "        response_mask: bool mask, True for response tokens (what we train on)\n",
+    "    \"\"\"\n",
+    "    def __init__(self, dataset, tokenizer, max_prompt_len, max_response_len):\n",
+    "        self.data = dataset\n",
+    "        self.tokenizer = tokenizer\n",
+    "        self.max_prompt_len = max_prompt_len\n",
+    "        self.max_response_len = max_response_len\n",
+    "        self.total_len = max_prompt_len + max_response_len\n",
+    "    \n",
+    "    def __len__(self):\n",
+    "        return len(self.data)\n",
+    "    \n",
+    "    def __getitem__(self, idx):\n",
+    "        item = self.data[idx]\n",
+    "        \n",
+    "        # Build prompt\n",
+    "        instruction = item['instruction']\n",
+    "        if item.get('input', ''):\n",
+    "            instruction = instruction + ' ' + item['input']\n",
+    "        response = item['output']\n",
+    "        \n",
+    "        # Tokenize separately\n",
+    "        prompt_tokens = [USER_TOKEN] + self.tokenizer.encode(instruction)[:self.max_prompt_len - 2] + [ASST_TOKEN]\n",
+    "        response_tokens = self.tokenizer.encode(response)[:self.max_response_len - 1] + [END_TOKEN]\n",
+    "        \n",
+    "        # Combine\n",
+    "        input_ids = prompt_tokens + response_tokens\n",
+    "        prompt_len = len(prompt_tokens)\n",
+    "        \n",
+    "        # Pad or truncate to fixed length\n",
+    "        if len(input_ids) < self.total_len:\n",
+    "            pad_len = self.total_len - len(input_ids)\n",
+    "            input_ids = input_ids + [tokenizer.eos_token_id] * pad_len\n",
+    "        else:\n",
+    "            input_ids = input_ids[:self.total_len]\n",
+    "        \n",
+    "        input_ids = torch.tensor(input_ids, dtype=torch.long)\n",
+    "        \n",
+    "        # Response mask: True for response positions only\n",
+    "        response_mask = torch.zeros(self.total_len, dtype=torch.bool)\n",
+    "        response_mask[prompt_len:prompt_len + len(response_tokens)] = True\n",
+    "        \n",
+    "        return input_ids, response_mask\n",
+    "\n",
+    "chat_dataset = ChatDataset(alpaca, tokenizer, ft_config.max_prompt_len, ft_config.max_response_len)\n",
+    "chat_loader = DataLoader(chat_dataset, batch_size=ft_config.ft_batch_size, shuffle=True, num_workers=2, pin_memory=True)\n",
+    "\n",
+    "# Test\n",
+    "test_ids, test_mask = chat_dataset[0]\n",
+    "print(f'\\nExample:')\n",
+    "print(f'Full sequence: {tokenizer.decode(test_ids[:40])}...')\n",
+    "print(f'Prompt tokens: {test_mask.sum().item()} response positions out of {len(test_ids)}')\n",
+    "print(f'\\nPrompt part: {tokenizer.decode(test_ids[~test_mask][:30])}')\n",
+    "print(f'Response part: {tokenizer.decode(test_ids[test_mask][:30])}')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "ft_train",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ============================================================\n",
+    "# FINE-TUNING LOOP\n",
+    "# ============================================================\n",
+    "\n",
+    "# Fresh optimizer with lower LR\n",
+    "ft_optimizer = torch.optim.AdamW(\n",
+    "    model_unwrapped.parameters(),\n",
+    "    lr=ft_config.ft_lr,\n",
+    "    betas=(0.9, 0.98),\n",
+    "    weight_decay=0.01,\n",
+    ")\n",
+    "ft_scaler = GradScaler('cuda')\n",
+    "ft_ema = EMA(model_unwrapped, decay=0.999)  # Faster EMA for fine-tuning\n",
+    "\n",
+    "model_unwrapped.train()\n",
+    "ft_losses = []\n",
+    "ft_accuracies = []\n",
+    "ft_start = time.time()\n",
+    "chat_iter = iter(chat_loader)\n",
+    "\n",
+    "print(f'Fine-tuning for {ft_config.ft_steps} steps...')\n",
+    "print(f'Batch size: {ft_config.ft_batch_size}')\n",
+    "print('=' * 60)\n",
+    "\n",
+    "for step in range(1, ft_config.ft_steps + 1):\n",
+    "    # LR schedule: linear warmup + cosine decay\n",
+    "    lr = get_lr(step, ft_config.ft_warmup, ft_config.ft_steps, ft_config.ft_lr)\n",
+    "    for pg in ft_optimizer.param_groups:\n",
+    "        pg['lr'] = lr\n",
+    "\n",
+    "    try:\n",
+    "        input_ids, response_mask = next(chat_iter)\n",
+    "    except StopIteration:\n",
+    "        chat_iter = iter(chat_loader)\n",
+    "        input_ids, response_mask = next(chat_iter)\n",
+    "\n",
+    "    input_ids = input_ids.to(device)\n",
+    "    response_mask = response_mask.to(device)\n",
+    "\n",
+    "    ft_optimizer.zero_grad()\n",
+    "\n",
+    "    with autocast('cuda', dtype=torch.float16):\n",
+    "        B, L = input_ids.shape\n",
+    "\n",
+    "        # Sample timestep\n",
+    "        t = model_unwrapped.noise_schedule.sample_t(B, device)\n",
+    "\n",
+    "        # Forward process: mask ONLY response tokens\n",
+    "        # Prompt tokens stay unmasked (model can always see them)\n",
+    "        alpha_t = model_unwrapped.noise_schedule.alpha(t)[:, None]  # [B, 1]\n",
+    "        mask_prob = 1.0 - alpha_t\n",
+    "        noise_mask = (torch.rand_like(input_ids.float()) < mask_prob) & response_mask\n",
+    "        z_t = torch.where(noise_mask, config.mask_token_id, input_ids)\n",
+    "\n",
+    "        # Forward pass\n",
+    "        hidden = model_unwrapped.forward_hidden(z_t, t)\n",
+    "\n",
+    "        # Loss only at masked response positions\n",
+    "        masked_hidden = hidden[noise_mask]\n",
+    "        masked_targets = input_ids[noise_mask]\n",
+    "\n",
+    "        if masked_hidden.shape[0] > 0:\n",
+    "            masked_logits = F.linear(masked_hidden, model_unwrapped.output_proj.weight)\n",
+    "            masked_logits[:, config.mask_token_id] = -1e9\n",
+    "            ce_loss = F.cross_entropy(masked_logits, masked_targets, reduction='none')\n",
+    "            weight = model_unwrapped.noise_schedule.loss_weight(t)\n",
+    "            weight_expanded = weight[:, None].expand(B, L)[noise_mask]\n",
+    "            loss = (ce_loss * weight_expanded).mean()\n",
+    "\n",
+    "            with torch.no_grad():\n",
+    "                acc = (masked_logits.argmax(-1) == masked_targets).float().mean().item()\n",
+    "        else:\n",
+    "            loss = torch.tensor(0.0, device=device)\n",
+    "            acc = 1.0\n",
+    "\n",
+    "    ft_scaler.scale(loss).backward()\n",
+    "    ft_scaler.unscale_(ft_optimizer)\n",
+    "    grad_norm = nn.utils.clip_grad_norm_(model_unwrapped.parameters(), 1.0)\n",
+    "    ft_scaler.step(ft_optimizer)\n",
+    "    ft_scaler.update()\n",
+    "    ft_ema.update(model_unwrapped)\n",
+    "\n",
+    "    ft_losses.append(loss.item())\n",
+    "    ft_accuracies.append(acc)\n",
+    "\n",
+    "    if step % ft_config.log_every == 0:\n",
+    "        elapsed = time.time() - ft_start\n",
+    "        avg_loss = np.mean(ft_losses[-ft_config.log_every:])\n",
+    "        avg_acc = np.mean(ft_accuracies[-ft_config.log_every:])\n",
+    "        eta = (ft_config.ft_steps - step) / (step / elapsed) / 60\n",
+    "        print(f'Step {step:>5d}/{ft_config.ft_steps} | Loss: {avg_loss:.4f} | Acc: {avg_acc:.3f} | LR: {lr:.2e} | Grad: {grad_norm:.2f} | ETA: {eta:.1f}m')\n",
+    "\n",
+    "    # Generate chat samples\n",
+    "    if step % ft_config.sample_every == 0:\n",
+    "        print(f\"\\n{'='*60}\")\n",
+    "        print(f'Chat samples at step {step}:')\n",
+    "        ft_ema.apply_shadow(model_unwrapped)\n",
+    "        model_unwrapped.eval()\n",
+    "\n",
+    "        test_prompts = [\n",
+    "            'What is the moon?',\n",
+    "            'Write a short poem about the ocean.',\n",
+    "            'Explain what a computer is.',\n",
+    "            'What is the meaning of life?',\n",
+    "        ]\n",
+    "\n",
+    "        for prompt in test_prompts:\n",
+    "            # Tokenize prompt\n",
+    "            prompt_tokens = [USER_TOKEN] + tokenizer.encode(prompt)[:ft_config.max_prompt_len - 2] + [ASST_TOKEN]\n",
+    "            prompt_len = len(prompt_tokens)\n",
+    "            total_len = prompt_len + ft_config.max_response_len\n",
+    "\n",
+    "            # Start with prompt + all masks for response\n",
+    "            x = torch.full((1, total_len), config.mask_token_id, dtype=torch.long, device=device)\n",
+    "            x[0, :prompt_len] = torch.tensor(prompt_tokens, dtype=torch.long, device=device)\n",
+    "\n",
+    "            # Diffusion sampling \u2014 only unmask response positions\n",
+    "            timesteps = torch.linspace(1.0 - 1e-5, 1e-5, 128 + 1, device=device)\n",
+    "            for i in range(128):\n",
+    "                t_now = timesteps[i]\n",
+    "                t_next = timesteps[i + 1]\n",
+    "                alpha_now = model_unwrapped.noise_schedule.alpha(t_now)\n",
+    "                alpha_next = model_unwrapped.noise_schedule.alpha(t_next)\n",
+    "\n",
+    "                t_batch = torch.full((1,), t_now.item(), device=device)\n",
+    "                logits = model_unwrapped.forward_full(x, t_batch)\n",
+    "                probs = F.softmax(logits / 0.7, dim=-1)\n",
+    "\n",
+    "                unmask_prob = ((alpha_next - alpha_now) / (1.0 - alpha_now + 1e-8)).clamp(0, 1)\n",
+    "                is_masked = (x == config.mask_token_id)\n",
+    "                unmask = is_masked & (torch.rand_like(x.float()) < unmask_prob)\n",
+    "\n",
+    "                if unmask.any():\n",
+    "                    flat_probs = probs.reshape(-1, config.vocab_size)\n",
+    "                    sampled = torch.multinomial(flat_probs, 1).reshape(1, total_len)\n",
+    "                    x = torch.where(unmask, sampled, x)\n",
+    "\n",
+    "            # Final cleanup\n",
+    "            is_masked = (x == config.mask_token_id)\n",
+    "            if is_masked.any():\n",
+    "                t_batch = torch.full((1,), 1e-5, device=device)\n",
+    "                logits = model_unwrapped.forward_full(x, t_batch)\n",
+    "                probs = F.softmax(logits / 0.7, dim=-1)\n",
+    "                flat_probs = probs.reshape(-1, config.vocab_size)\n",
+    "                sampled = torch.multinomial(flat_probs, 1).reshape(1, total_len)\n",
+    "                x = torch.where(is_masked, sampled, x)\n",
+    "\n",
+    "            # Decode response only\n",
+    "            response_tokens = x[0, prompt_len:].cpu().tolist()\n",
+    "            # Cut at END token\n",
+    "            if END_TOKEN in response_tokens:\n",
+    "                response_tokens = response_tokens[:response_tokens.index(END_TOKEN)]\n",
+    "            response = tokenizer.decode(response_tokens, skip_special_tokens=True)\n",
+    "            print(f'\\n  User: {prompt}')\n",
+    "            print(f'  Bot:  {response}')\n",
+    "\n",
+    "        model_unwrapped.train()\n",
+    "        ft_ema.restore(model_unwrapped)\n",
+    "        print(f\"{'='*60}\\n\")\n",
+    "\n",
+    "# Save fine-tuned model\n",
+    "torch.save({\n",
+    "    'step': step,\n",
+    "    'model_state_dict': model_unwrapped.state_dict(),\n",
+    "    'ema_shadow': ft_ema.shadow,\n",
+    "    'config': config,\n",
+    "}, 'checkpoint_chat.pt')\n",
+    "print('Fine-tuning complete! Saved checkpoint_chat.pt')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "chat_header",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Chat with your Diffusion LM\n",
+    "\n",
+    "Type a message and watch the response **materialize from noise** via the diffusion process."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "chat_interface",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ============================================================\n",
+    "# CHAT INTERFACE WITH DIFFUSION VISUALIZATION\n",
+    "# ============================================================\n",
+    "\n",
+    "from IPython.display import clear_output, display\n",
+    "import time as _time\n",
+    "\n",
+    "# Load EMA weights\n",
+    "ft_ema.apply_shadow(model_unwrapped)\n",
+    "model_unwrapped.eval()\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def chat(prompt: str, steps: int = 64, temperature: float = 0.7, show_diffusion: bool = True):\n",
+    "    \"\"\"Chat with the diffusion model.\n",
+    "    \n",
+    "    Args:\n",
+    "        prompt: Your message\n",
+    "        steps: Denoising steps (more = better quality, slower)\n",
+    "        temperature: Sampling temperature (lower = more focused)\n",
+    "        show_diffusion: Show the step-by-step unmasking process\n",
+    "    \"\"\"\n",
+    "    # Tokenize prompt\n",
+    "    prompt_tokens = [USER_TOKEN] + tokenizer.encode(prompt)[:ft_config.max_prompt_len - 2] + [ASST_TOKEN]\n",
+    "    prompt_len = len(prompt_tokens)\n",
+    "    total_len = prompt_len + ft_config.max_response_len\n",
+    "\n",
+    "    # Initialize: prompt (visible) + all masks (response)\n",
+    "    x = torch.full((1, total_len), config.mask_token_id, dtype=torch.long, device=device)\n",
+    "    x[0, :prompt_len] = torch.tensor(prompt_tokens, dtype=torch.long, device=device)\n",
+    "\n",
+    "    timesteps_sched = torch.linspace(1.0 - 1e-5, 1e-5, steps + 1, device=device)\n",
+    "    snapshot_steps = set([int(steps * p) for p in [0, 0.1, 0.2, 0.35, 0.5, 0.7, 0.85, 1.0]])\n",
+    "\n",
+    "    if show_diffusion:\n",
+    "        print(f'User: {prompt}')\n",
+    "        print(f'\\n--- Diffusion Process ({steps} steps) ---\\n')\n",
+    "\n",
+    "    for i in range(steps):\n",
+    "        t_now = timesteps_sched[i]\n",
+    "        t_next = timesteps_sched[i + 1]\n",
+    "        alpha_now = model_unwrapped.noise_schedule.alpha(t_now)\n",
+    "        alpha_next = model_unwrapped.noise_schedule.alpha(t_next)\n",
+    "\n",
+    "        t_batch = torch.full((1,), t_now.item(), device=device)\n",
+    "        logits = model_unwrapped.forward_full(x, t_batch)\n",
+    "        probs = F.softmax(logits / temperature, dim=-1)\n",
+    "\n",
+    "        unmask_prob = ((alpha_next - alpha_now) / (1.0 - alpha_now + 1e-8)).clamp(0, 1)\n",
+    "        is_masked = (x == config.mask_token_id)\n",
+    "        unmask = is_masked & (torch.rand_like(x.float()) < unmask_prob)\n",
+    "\n",
+    "        if unmask.any():\n",
+    "            flat_probs = probs.reshape(-1, config.vocab_size)\n",
+    "            sampled = torch.multinomial(flat_probs, 1).reshape(1, total_len)\n",
+    "            x = torch.where(unmask, sampled, x)\n",
+    "\n",
+    "        # Show snapshot\n",
+    "        if show_diffusion and i in snapshot_steps:\n",
+    "            resp_tokens = x[0, prompt_len:].cpu().tolist()\n",
+    "            text = ''\n",
+    "            for tok in resp_tokens:\n",
+    "                if tok == config.mask_token_id:\n",
+    "                    text += ' \\u2588'\n",
+    "                elif tok == END_TOKEN:\n",
+    "                    break\n",
+    "                else:\n",
+    "                    text += tokenizer.decode([tok])\n",
+    "            pct = (1 - is_masked[:, prompt_len:].float().mean()).item() * 100\n",
+    "            print(f'  [{pct:5.1f}% revealed] {text[:200]}')\n",
+    "\n",
+    "    # Final cleanup\n",
+    "    is_masked = (x == config.mask_token_id)\n",
+    "    if is_masked.any():\n",
+    "        t_batch = torch.full((1,), 1e-5, device=device)\n",
+    "        logits = model_unwrapped.forward_full(x, t_batch)\n",
+    "        probs = F.softmax(logits / temperature, dim=-1)\n",
+    "        flat_probs = probs.reshape(-1, config.vocab_size)\n",
+    "        sampled = torch.multinomial(flat_probs, 1).reshape(1, total_len)\n",
+    "        x = torch.where(is_masked, sampled, x)\n",
+    "\n",
+    "    # Decode final response\n",
+    "    response_tokens = x[0, prompt_len:].cpu().tolist()\n",
+    "    if END_TOKEN in response_tokens:\n",
+    "        response_tokens = response_tokens[:response_tokens.index(END_TOKEN)]\n",
+    "    response = tokenizer.decode(response_tokens, skip_special_tokens=True)\n",
+    "\n",
+    "    if show_diffusion:\n",
+    "        print(f'\\n--- Final ---')\n",
+    "    print(f'\\nUser: {prompt}')\n",
+    "    print(f'Bot:  {response}')\n",
+    "    return response\n",
+    "\n",
+    "print('Chat function ready! Usage: chat(\"your message here\")')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "chat_examples",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Try it out!\n",
+    "chat('What is the moon?')\n",
+    "print('\\n' + '='*60 + '\\n')\n",
+    "chat('Write a short poem about the ocean.')\n",
+    "print('\\n' + '='*60 + '\\n')\n",
+    "chat('Explain what a computer is to a child.')\n",
+    "print('\\n' + '='*60 + '\\n')\n",
+    "chat('What are three things that make people happy?')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "ft_upload",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Upload fine-tuned model to HuggingFace\n",
+    "from huggingface_hub import HfApi\n",
+    "TOKEN = 'YOUR_HF_TOKEN_HERE'\n",
+    "api = HfApi(token=TOKEN)\n",
+    "\n",
+    "api.upload_file(\n",
+    "    path_or_fileobj='checkpoint_chat.pt',\n",
+    "    path_in_repo='checkpoint_chat.pt',\n",
+    "    repo_id='chipling/opium-mdlm',\n",
+    "    repo_type='model',\n",
+    "    token=TOKEN,\n",
+    ")\n",
+    "print('Chat model uploaded to HuggingFace!')\n"
+   ]
   }
  ],
  "metadata": {