{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "kaggle": {
      "accelerator": "gpu",
      "dataSources": [],
      "isInternetEnabled": true,
      "language": "python",
      "sourceType": "notebook"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# 🤖 MCP-Agent-1.7B — Training on Kaggle (FREE GPU)\n",
        "\n",
        "**What:** Fine-tune Qwen3-1.7B to natively speak Model Context Protocol (MCP) for tool calling.\n",
        "\n",
        "**GPU:** Kaggle P100/T4 (16GB VRAM) — **completely free**, 30 hrs/week\n",
        "\n",
        "**Time:** ~2 hours\n",
        "\n",
        "---\n",
        "\n",
        "⚡ **Before running:** Go to **Settings (right panel) → Accelerator → GPU P100 or GPU T4×2**\n",
        "\n",
        "⚡ **Also:** Settings → Internet → **Turn ON** (needed to download model & push to Hub)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# ============================================================\n",
        "# CELL 1: Check GPU & Install Dependencies\n",
        "# ============================================================\n",
        "!nvidia-smi\n",
        "\n",
        "import torch\n",
        "assert torch.cuda.is_available(), \"❌ No GPU! Go to Settings → Accelerator → GPU\"\n",
        "print(f\"\\n✅ GPU: {torch.cuda.get_device_name(0)}\")\n",
        "print(f\"✅ VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB\")\n",
        "\n",
        "!pip install -q transformers trl peft datasets accelerate bitsandbytes huggingface_hub\n",
        "print(\"\\n✅ All packages installed!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# ============================================================\n",
        "# CELL 2: Login to HuggingFace\n",
        "# ============================================================\n",
        "# Get your token at: https://huggingface.co/settings/tokens (needs WRITE permission)\n",
        "\n",
        "from huggingface_hub import notebook_login\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# ============================================================\n",
        "# CELL 3: Load Everything & Start Training\n",
        "# ============================================================\n",
        "# This single cell does it all — loads data, configures LoRA,\n",
        "# sets up training, and runs for 3 epochs (~2 hours).\n",
        "# Don't close the tab while it runs!\n",
        "# ============================================================\n",
        "\n",
        "import os, torch\n",
        "from datasets import load_dataset\n",
        "from peft import LoraConfig\n",
        "from trl import SFTConfig, SFTTrainer\n",
        "from transformers import AutoTokenizer\n",
        "\n",
        "# --- Load dataset (16.5K MCP tool-calling conversations) ---\n",
        "print(\"📦 Loading dataset...\")\n",
        "dataset = load_dataset(\"muhammadtlha944/mcp-agent-training-data\")\n",
        "print(f\"   Train: {len(dataset['train']):,} | Val: {len(dataset['validation']):,}\")\n",
        "\n",
        "# --- Load tokenizer ---\n",
        "print(\"📦 Loading tokenizer...\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen3-1.7B\", trust_remote_code=True)\n",
        "\n",
        "# --- LoRA config (only train ~2% of parameters) ---\n",
        "peft_config = LoraConfig(\n",
        "    r=16,                        # Rank of adapter matrices\n",
        "    lora_alpha=32,               # Scaling factor (2x rank)\n",
        "    lora_dropout=0.05,           # Regularization\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\",\n",
        "    target_modules=\"all-linear\", # Apply LoRA to ALL linear layers\n",
        ")\n",
        "\n",
        "# --- Detect GPU type for precision ---\n",
        "gpu_name = torch.cuda.get_device_name(0).lower()\n",
        "use_bf16 = \"a100\" in gpu_name or \"a10\" in gpu_name or \"l4\" in gpu_name or \"h100\" in gpu_name\n",
        "print(f\"   Using {'bf16' if use_bf16 else 'fp16'} precision for {torch.cuda.get_device_name(0)}\")\n",
        "\n",
        "# --- Training config ---\n",
        "training_args = SFTConfig(\n",
        "    output_dir=\"/kaggle/working/mcp-agent-checkpoints\",\n",
        "    \n",
        "    # Core hyperparameters\n",
        "    num_train_epochs=3,\n",
        "    per_device_train_batch_size=4,\n",
        "    gradient_accumulation_steps=4,       # Effective batch = 16\n",
        "    learning_rate=2e-4,                  # 10x base LR for LoRA\n",
        "    weight_decay=0.01,\n",
        "    lr_scheduler_type=\"cosine\",\n",
        "    warmup_ratio=0.1,\n",
        "    max_grad_norm=1.0,\n",
        "    max_seq_length=2048,\n",
        "    \n",
        "    # Memory optimization\n",
        "    bf16=use_bf16,\n",
        "    fp16=not use_bf16,\n",
        "    gradient_checkpointing=True,\n",
        "    gradient_checkpointing_kwargs={\"use_reentrant\": False},\n",
        "    \n",
        "    # Logging\n",
        "    logging_steps=10,\n",
        "    logging_first_step=True,\n",
        "    logging_strategy=\"steps\",\n",
        "    \n",
        "    # Evaluation & Checkpoints\n",
        "    eval_strategy=\"steps\",\n",
        "    eval_steps=200,\n",
        "    per_device_eval_batch_size=4,\n",
        "    save_strategy=\"steps\",\n",
        "    save_steps=200,\n",
        "    save_total_limit=2,\n",
        "    load_best_model_at_end=True,\n",
        "    metric_for_best_model=\"eval_loss\",\n",
        "    \n",
        "    # Push to Hub\n",
        "    push_to_hub=True,\n",
        "    hub_model_id=\"muhammadtlha944/MCP-Agent-1.7B\",\n",
        "    hub_strategy=\"end\",\n",
        "    \n",
        "    # Misc\n",
        "    seed=42,\n",
        "    dataloader_num_workers=2,\n",
        "    optim=\"adamw_torch\",\n",
        ")\n",
        "\n",
        "# --- Create trainer ---\n",
        "print(\"\\n🔧 Loading Qwen3-1.7B and applying LoRA...\")\n",
        "trainer = SFTTrainer(\n",
        "    model=\"Qwen/Qwen3-1.7B\",\n",
        "    args=training_args,\n",
        "    train_dataset=dataset[\"train\"],\n",
        "    eval_dataset=dataset[\"validation\"],\n",
        "    peft_config=peft_config,\n",
        "    processing_class=tokenizer,\n",
        ")\n",
        "\n",
        "trainable = sum(p.numel() for p in trainer.model.parameters() if p.requires_grad)\n",
        "total = sum(p.numel() for p in trainer.model.parameters())\n",
        "print(f\"   Total params: {total:,}\")\n",
        "print(f\"   Trainable (LoRA): {trainable:,} ({100*trainable/total:.2f}%)\")\n",
        "print(f\"   GPU memory: {torch.cuda.memory_allocated()/1e9:.1f} GB\")\n",
        "\n",
        "# --- TRAIN! ---\n",
        "print(f\"\\n🚀 Starting training (~2 hours)...\\n\")\n",
        "train_result = trainer.train()\n",
        "\n",
        "print(f\"\\n✅ Training done! Loss: {train_result.metrics.get('train_loss', 'N/A')}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# ============================================================\n",
        "# CELL 4: Evaluate & Push to HuggingFace Hub\n",
        "# ============================================================\n",
        "\n",
        "# Final eval\n",
        "print(\"📊 Evaluating...\")\n",
        "eval_metrics = trainer.evaluate()\n",
        "print(f\"   Eval loss: {eval_metrics['eval_loss']:.4f}\")\n",
        "\n",
        "# Save metrics\n",
        "trainer.log_metrics(\"train\", train_result.metrics)\n",
        "trainer.save_metrics(\"train\", train_result.metrics)\n",
        "trainer.log_metrics(\"eval\", eval_metrics)\n",
        "trainer.save_metrics(\"eval\", eval_metrics)\n",
        "\n",
        "# Push\n",
        "print(\"\\n🚀 Pushing to HuggingFace Hub...\")\n",
        "trainer.push_to_hub(\n",
        "    commit_message=\"MCP-Agent-1.7B: LoRA fine-tuned Qwen3-1.7B for MCP tool calling\",\n",
        "    tags=[\"mcp\", \"tool-calling\", \"function-calling\", \"agent\", \"qwen3\", \"lora\"],\n",
        ")\n",
        "\n",
        "print(f\"\\n{'='*60}\")\n",
        "print(f\"🎉 MCP-Agent-1.7B is LIVE!\")\n",
        "print(f\"{'='*60}\")\n",
        "print(f\"📦 https://huggingface.co/muhammadtlha944/MCP-Agent-1.7B\")\n",
        "print(f\"📊 Train loss: {train_result.metrics.get('train_loss', 'N/A')}\")\n",
        "print(f\"📊 Eval loss: {eval_metrics['eval_loss']:.4f}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# ============================================================\n",
        "# CELL 5: Test the Model!\n",
        "# ============================================================\n",
        "from transformers import pipeline\n",
        "\n",
        "pipe = pipeline(\"text-generation\", model=trainer.model, tokenizer=tokenizer,\n",
        "                max_new_tokens=512, do_sample=True, temperature=0.7)\n",
        "\n",
        "tests = [\n",
        "    \"Find all Python files in src/ that import pandas\",\n",
        "    \"Clone the repo, find all TODO comments, create a summary\",\n",
        "    \"Delete the database\",\n",
        "]\n",
        "\n",
        "for i, prompt in enumerate(tests, 1):\n",
        "    messages = [\n",
        "        {\"role\": \"system\", \"content\": \"You are an MCP agent with tools: github_search, read_file, shell_exec, sqlite_query. Use JSON-RPC format. Ask for clarification when needed. Refuse dangerous requests.\"},\n",
        "        {\"role\": \"user\", \"content\": prompt}\n",
        "    ]\n",
        "    result = pipe({\"messages\": messages})\n",
        "    response = result[0]['generated_text'][-1]['content']\n",
        "    print(f\"\\n{'='*50}\")\n",
        "    print(f\"TEST {i}: {prompt}\")\n",
        "    print(f\"{'='*50}\")\n",
        "    print(f\"🤖 {response}\")"
      ]
    }
  ]
}