sft/proxy_tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "32818444-ac19-48f3-b9fb-6f9dd8137d7e",
   "metadata": {},
   "source": [
    "### 确保使用具有 GPU 的计算机"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9f601b92-5f22-4361-a9a1-dd58996728c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of GPUs available: 1\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    num_gpus = torch.cuda.device_count()\n",
    "    print(f\"Number of GPUs available: {num_gpus}\")\n",
    "    if num_gpus < 1:\n",
    "        raise Exception(\"At least 1 GPU is required for this setup.\")\n",
    "else:\n",
    "    raise Exception(\"No GPUs found, make sure a GPU is available\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13606944-5e26-49a1-af13-b5e1d9b37ce6",
   "metadata": {},
   "source": [
    "### 使用未修改的模型生成文本\n",
    "\n",
    "请注意 `Llama-2-13b-chat-hf` 是一个对照实验。我们的目标是使用更小的可用 `Llama-2-7b-chat-hf` 模型，以便在将其应用于 `Llama-2-13b-hf` 基本模型时实现相同的性能。\n",
    "\n",
    "接下来，让我们检查一下这些模型在不应用代理调优的情况下产生的输出："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6c2903c0-a96a-4665-9899-400f8c14513c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading ./autodl-tmp/Llama-2-7b-hf\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c167a114f94548989ec5a14d2e4d4ae1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "\n",
      "Asking ./autodl-tmp/Llama-2-7b-hf:\n",
      "\u001b[5;34;46m \n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I find 3 more.\n",
      "I have 10 apples.\n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I find 3 more.\n",
      "I have 10 apples.\n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I find 3 more.\n",
      "I have 10 apples.\n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I find 3 more.\n",
      "I have 10 apples.\n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I find 3 more.\n",
      "I have 10 apples.\n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I find 3 more.\n",
      "I have 10 apples.\n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I find 3 more.\n",
      "I have 10 apples.\n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I find 3 more.\n",
      "I have 10 apples.\n",
      "I have 5 apples.\n",
      "I eat 2.\n",
      "I \u001b[0m \n",
      "Length: 115\n",
      "Loading ./autodl-tmp/Llama-2-7b-chat-hf\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd11b89dccd54ed388b2d81643cd5eec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "\n",
      "Asking ./autodl-tmp/Llama-2-7b-chat-hf:\n",
      "\u001b[5;34;46m \n",
      "\n",
      "Answer: You have 5 apples in total.\n",
      "\n",
      "Explanation:\n",
      "\n",
      "When you started your journey, you had 5 apples.\n",
      "\n",
      "When you ate 2 apples, you had 5 - 2 = 3 apples left.\n",
      "\n",
      "When you found 3 more apples on your way home, you now have 3 + 3 = 6 apples in total.\n",
      "\n",
      "So, the final answer is 6 apples. \u001b[0m \n",
      "Length: 75\n",
      "Loading ./autodl-tmp/Llama-2-13b-hf\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a59eb5288ed4530a16a38366e939238",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "\n",
      "Asking ./autodl-tmp/Llama-2-13b-hf:\n",
      "\u001b[5;34;46m \n",
      "If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have?\n",
      "If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have? If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have? If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have? If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have? If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have? If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have? If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have? If I have 5 apples and eat 2, but then find 3 more on my way home \u001b[0m \n",
      "Length: 213\n",
      "Loading ./autodl-tmp/Llama-2-13b-chat-hf\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9f2f2b444fc4a44a8ec5e8c7f6e6050",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "\n",
      "Asking ./autodl-tmp/Llama-2-13b-chat-hf:\n",
      "\u001b[5;34;46m \n",
      "\n",
      "You start with 5 apples and eat 2, so you have 5 - 2 = 3 apples left.\n",
      "Then, you find 3 more apples on your way home, so you have 3 + 3 = 6 apples in total. \u001b[0m \n",
      "Length: 59\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from transformers import BitsAndBytesConfig\n",
    "import os\n",
    "import torch\n",
    "\n",
    "device = \"cuda\"\n",
    "\n",
    "def load_model_and_tokenizer(model_dir):\n",
    "    bnb_config = BitsAndBytesConfig(  \n",
    "        load_in_4bit= True,\n",
    "        bnb_4bit_quant_type= \"nf4\",\n",
    "        bnb_4bit_compute_dtype= torch.bfloat16,\n",
    "        bnb_4bit_use_double_quant= False,\n",
    "    )\n",
    "    \n",
    "    if \"7b\" in model_name:\n",
    "        model = AutoModelForCausalLM.from_pretrained(\n",
    "            model_name, \n",
    "            device_map=device,\n",
    "            #attn_implementation=\"flash_attention_2\", \n",
    "            torch_dtype=torch.bfloat16\n",
    "        )\n",
    "    else:\n",
    "        model = AutoModelForCausalLM.from_pretrained(\n",
    "            model_name, \n",
    "            device_map=device,\n",
    "            #attn_implementation=\"flash_attention_2\", \n",
    "            torch_dtype=torch.bfloat16, \n",
    "            #load_in_4bit=True,\n",
    "            quantization_config=bnb_config,\n",
    "        )\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "    return model, tokenizer\n",
    "\n",
    "\n",
    "def generate_text(model, tokenizer, input_text, max_length):\n",
    "    device = next(model.parameters()).device\n",
    "    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)\n",
    "    generated_tokens = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_length):\n",
    "            \n",
    "            outputs = model(input_ids)\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.softmax(logits[:, -1, :], dim=-1)\n",
    "            next_token_id = torch.argmax(predictions).unsqueeze(0)\n",
    "            generated_tokens.append(next_token_id.item())\n",
    "            \n",
    "            # Append the new token to the input sequence for the next iteration\n",
    "            input_ids = torch.cat([input_ids, next_token_id.unsqueeze(0)], dim=1)\n",
    "\n",
    "            if next_token_id.item() == tokenizer.eos_token_id:\n",
    "                break\n",
    "\n",
    "    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)\n",
    "\n",
    "    return generated_text\n",
    "\n",
    "\n",
    "model_names = [\n",
    "    \"./autodl-tmp/Llama-2-7b-hf\",\n",
    "    \"./autodl-tmp/Llama-2-7b-chat-hf\",\n",
    "    \"./autodl-tmp/Llama-2-13b-hf\",\n",
    "    \"./autodl-tmp/Llama-2-13b-chat-hf\",\n",
    "]\n",
    "\n",
    "for model_name in model_names:\n",
    "    input_text = (\"If I have 5 apples and eat 2, but then find 3 more on my way home, how many do I have?\")\n",
    "\n",
    "    print(\"Loading\", model_name)\n",
    "    model, tokenizer = load_model_and_tokenizer(model_name)\n",
    "    \n",
    "    print(f\"{50*'='}\\n\\nAsking {model_name}:\")\n",
    "    generated = generate_text(model, tokenizer, input_text, max_length=256)\n",
    "    print(\"\\033[5;34;46m\", generated, \"\\033[0m\", f\"\\nLength: {len((input_text+generated).split(' '))}\")\n",
    "\n",
    "    del model\n",
    "    del tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5cb46fb-853f-467f-91a7-4a095e608caf",
   "metadata": {},
   "source": [
    "正如我们所看到的，只有 `Llama-2-13b-chat-hf` 模型（我们的对照实验）才能获得正确的响应。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a155e169-fdbd-4023-938b-902336cee394",
   "metadata": {},
   "source": [
    "### 使用 Llama 2 7B Chat 代理调整 13B Base 模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "79827dbb-aee6-43d6-acb6-88aabcd001d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b7e8236511cd4edd86a6ea67befb0a21",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   0%|          | 0/256 [00:00<?, ?it/s]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "efb06f7cd8734ad985a47f695f9d2554",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7ca18403d895428daf92974d94a2e901",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c12bdd94c474471ae6c2efc6e84d881",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   0%|          | 1/256 [00:27<1:57:12, 27.58s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27cf259230514cdea595fc801a9979c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "73c111fb9e894cc8be1aa5eeb3cd6586",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a24378cad39b4c86a651e194c591a3fa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   1%|          | 2/256 [00:55<1:56:30, 27.52s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c617231e2b24cde9e2d2ea6b5aec814",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3d5bbaaa59bf4bdfbb491835c2467841",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99061d37d07f4782af90a32fd9ad2b01",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   1%|          | 3/256 [01:22<1:55:53, 27.48s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b7440c8a23ef4a3b94e7e4d2a77f88aa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c0f96c44538d4a4288a231800a1f4e2f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d107eb208b1b46b7bb82cc6091e2e76c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   2%|▏         | 4/256 [01:52<1:59:44, 28.51s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d0f593a32d54e9d9a4363555c96eb14",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "808e8fe1c8ae46b19f14c79f61494e4c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "55bd7c28523340f8b125763d28411bd6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   2%|▏         | 5/256 [02:22<2:00:43, 28.86s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fab3d098bce40d4b2062323131ea83c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b86f4a2597de4f0dae14a3b28550b959",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5590d22cdd1c46f5850ab6ce86566031",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   2%|▏         | 6/256 [02:51<2:01:20, 29.12s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95820d8dd98b4da48615c69062d4a6c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "666baba57d574e5db3f5cfe4c69c444e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a08049ad7fda4b149733d6028448fc41",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   3%|▎         | 7/256 [03:21<2:01:51, 29.36s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4978ec17fd174e5ab1caa1df3fa012c6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca564d577cff42cf85bcd9ecfd05775b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6343307c8a974d76a6301cd946fcb989",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   3%|▎         | 8/256 [03:51<2:01:58, 29.51s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "46345309807c45459f2962121a5602b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "13b8e48b92c54ed9897015196dcec509",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9cec97f9065f4d67a48f82522ff27e6b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   4%|▎         | 9/256 [04:21<2:02:35, 29.78s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42cd9f1b0413483ba09f983c6dd7eca6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fbcc52ff3fc54ed3906a01f15021c950",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5ae80833f5142cf8c5be44b31ef6bc3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   4%|▍         | 10/256 [04:54<2:05:58, 30.72s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "910db8068c65472e83bf62952d93ee93",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ccaea1f330e43aead751b2e4c00a390",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b484cdd80f974f6389f25438d333052a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   4%|▍         | 11/256 [05:35<2:17:39, 33.71s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c158e571872a4a47b7d757d1fd65a9bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2e399678ef904a24bb0d36091e7c1f63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa8d630f448e402388e120ee5125405e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   5%|▍         | 12/256 [06:16<2:26:15, 35.97s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e5cbb0e4642472b9832668264856568",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb7ed95dcf8a4cb393fbf9d786a0f74c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d2aef7ba5e7648889fa71309eade20eb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   5%|▌         | 13/256 [06:57<2:32:05, 37.55s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea188117be3946d6983315ce3de9452a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38705be7231240aca74d7a81fa01e57e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee481a279d9048448dcd90fb9b223a86",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   5%|▌         | 14/256 [07:38<2:35:55, 38.66s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d2d6edb02e64381ad9c8aff6c9b8a1c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76ce2428ff40431da9d1beb5e3e5dbb8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6fc60882ded5426c9d3bf75aa70492df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   6%|▌         | 15/256 [08:19<2:38:17, 39.41s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b3981b6d3d464f49a34b792faf3e21a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "70fd71a04344424bbb7af5b4d4f379ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0cdb96120c38463c89dfe6bc1cda98ed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   6%|▋         | 16/256 [09:00<2:39:42, 39.93s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "598ce174e99e443f8e4c684a3e5120e4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbe9773e21a54e10bf704601b808d948",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "47013309d6864ce1aa98317e74514ba8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   7%|▋         | 17/256 [09:39<2:36:58, 39.41s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1709e7cc99a543edb994834ec1914e15",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4af85fa01943436986fde4ec473edcaa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e972a34619c3427bb598418236159581",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   7%|▋         | 18/256 [10:16<2:33:32, 38.71s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af15e72d219f4727a9b139820c9e3c3b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e169c2a67cf4c6aa92455d848af4483",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ba64d819e8e467c994b47814c23719f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   7%|▋         | 19/256 [10:56<2:35:09, 39.28s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d275a819d9cd4ed6b4799bbe9a238c07",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a61ccc95027548e4b5a3be5d89af7b45",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5835fff7eda447c69bc527c5155ac94a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   8%|▊         | 20/256 [11:36<2:34:41, 39.33s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db0bc0d4dac54f1eb5cc588591a1712a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "673dd249a1d143c499d7a973a1a96c3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90724adc60204ac798974836b3720e78",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   8%|▊         | 21/256 [12:12<2:30:49, 38.51s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0a8f905c168409890e32fbd4247323e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4678f6eb65d443ca96126b59fd756fe1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a3ace196d0514e9b9c2232fec4c9a235",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating tokens:   8%|▊         | 21/256 [12:39<2:21:39, 36.17s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 65\u001b[0m\n\u001b[1;32m     61\u001b[0m     generated_text \u001b[38;5;241m=\u001b[39m tokenizer\u001b[38;5;241m.\u001b[39mdecode(generated_tokens, skip_special_tokens\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m     62\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m generated_text\n\u001b[0;32m---> 65\u001b[0m generated \u001b[38;5;241m=\u001b[39m \u001b[43mgenerate_proxy_tuning\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     66\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmodel_base_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m./autodl-tmp/Llama-2-7b-hf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     67\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmodel_tuned_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m./autodl-tmp/Llama-2-7b-chat-hf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     68\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmodel_target_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m./autodl-tmp/Llama-2-13b-hf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     69\u001b[0m \u001b[43m    \u001b[49m\u001b[43minput_text\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_text\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     70\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmax_length\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m256\u001b[39;49m\n\u001b[1;32m     71\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m     73\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\033\u001b[39;00m\u001b[38;5;124m[5;34;46m\u001b[39m\u001b[38;5;124m\"\u001b[39m, generated, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\033\u001b[39;00m\u001b[38;5;124m[0m\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[0;32mIn[5], line 42\u001b[0m, in \u001b[0;36mgenerate_proxy_tuning\u001b[0;34m(model_base_name, model_tuned_name, model_target_name, input_text, max_length)\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m model_tuned\n\u001b[1;32m     41\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m suppress_output():\n\u001b[0;32m---> 42\u001b[0m     model_target, _ \u001b[38;5;241m=\u001b[39m \u001b[43mload_model_and_tokenizer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel_target_name\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     43\u001b[0m logits_target \u001b[38;5;241m=\u001b[39m model_target(input_ids)\u001b[38;5;241m.\u001b[39mlogits\n\u001b[1;32m     44\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m model_target\n",
      "Cell \u001b[0;32mIn[4], line 24\u001b[0m, in \u001b[0;36mload_model_and_tokenizer\u001b[0;34m(model_dir)\u001b[0m\n\u001b[1;32m     17\u001b[0m     model \u001b[38;5;241m=\u001b[39m AutoModelForCausalLM\u001b[38;5;241m.\u001b[39mfrom_pretrained(\n\u001b[1;32m     18\u001b[0m         model_name, \n\u001b[1;32m     19\u001b[0m         device_map\u001b[38;5;241m=\u001b[39mdevice,\n\u001b[1;32m     20\u001b[0m         \u001b[38;5;66;03m#attn_implementation=\"flash_attention_2\", \u001b[39;00m\n\u001b[1;32m     21\u001b[0m         torch_dtype\u001b[38;5;241m=\u001b[39mtorch\u001b[38;5;241m.\u001b[39mbfloat16\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m---> 24\u001b[0m     model \u001b[38;5;241m=\u001b[39m \u001b[43mAutoModelForCausalLM\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     25\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmodel_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m     26\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdevice_map\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     27\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;66;43;03m#attn_implementation=\"flash_attention_2\", \u001b[39;49;00m\n\u001b[1;32m     28\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtorch_dtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbfloat16\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m     29\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;66;43;03m#load_in_4bit=True,\u001b[39;49;00m\n\u001b[1;32m     30\u001b[0m \u001b[43m        \u001b[49m\u001b[43mquantization_config\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbnb_config\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     31\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     32\u001b[0m tokenizer \u001b[38;5;241m=\u001b[39m AutoTokenizer\u001b[38;5;241m.\u001b[39mfrom_pretrained(model_name)\n\u001b[1;32m     33\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m model, tokenizer\n",
      "File \u001b[0;32m~/miniconda3/lib/python3.12/site-packages/transformers/models/auto/auto_factory.py:564\u001b[0m, in \u001b[0;36m_BaseAutoModelClass.from_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, *model_args, **kwargs)\u001b[0m\n\u001b[1;32m    562\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(config) \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m_model_mapping\u001b[38;5;241m.\u001b[39mkeys():\n\u001b[1;32m    563\u001b[0m     model_class \u001b[38;5;241m=\u001b[39m _get_model_class(config, \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m_model_mapping)\n\u001b[0;32m--> 564\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmodel_class\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    565\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpretrained_model_name_or_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mmodel_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mhub_kwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m    566\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    567\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    568\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUnrecognized configuration class \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mconfig\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m for this kind of AutoModel: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    569\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mModel type should be one of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(c\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mfor\u001b[39;00m\u001b[38;5;250m \u001b[39mc\u001b[38;5;250m \u001b[39m\u001b[38;5;129;01min\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m_model_mapping\u001b[38;5;241m.\u001b[39mkeys())\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    570\u001b[0m )\n",
      "File \u001b[0;32m~/miniconda3/lib/python3.12/site-packages/transformers/modeling_utils.py:4014\u001b[0m, in \u001b[0;36mPreTrainedModel.from_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, use_safetensors, *model_args, **kwargs)\u001b[0m\n\u001b[1;32m   4004\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m dtype_orig \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[1;32m   4005\u001b[0m         torch\u001b[38;5;241m.\u001b[39mset_default_dtype(dtype_orig)\n\u001b[1;32m   4007\u001b[0m     (\n\u001b[1;32m   4008\u001b[0m         model,\n\u001b[1;32m   4009\u001b[0m         missing_keys,\n\u001b[1;32m   4010\u001b[0m         unexpected_keys,\n\u001b[1;32m   4011\u001b[0m         mismatched_keys,\n\u001b[1;32m   4012\u001b[0m         offload_index,\n\u001b[1;32m   4013\u001b[0m         error_msgs,\n\u001b[0;32m-> 4014\u001b[0m     ) \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_load_pretrained_model\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   4015\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4016\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstate_dict\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4017\u001b[0m \u001b[43m        \u001b[49m\u001b[43mloaded_state_dict_keys\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# XXX: rename?\u001b[39;49;00m\n\u001b[1;32m   4018\u001b[0m \u001b[43m        \u001b[49m\u001b[43mresolved_archive_file\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4019\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpretrained_model_name_or_path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4020\u001b[0m \u001b[43m        \u001b[49m\u001b[43mignore_mismatched_sizes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mignore_mismatched_sizes\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4021\u001b[0m \u001b[43m        \u001b[49m\u001b[43msharded_metadata\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msharded_metadata\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4022\u001b[0m \u001b[43m        \u001b[49m\u001b[43m_fast_init\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_fast_init\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4023\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlow_cpu_mem_usage\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlow_cpu_mem_usage\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4024\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdevice_map\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice_map\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4025\u001b[0m \u001b[43m        \u001b[49m\u001b[43moffload_folder\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moffload_folder\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4026\u001b[0m \u001b[43m        \u001b[49m\u001b[43moffload_state_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moffload_state_dict\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4027\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtorch_dtype\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4028\u001b[0m \u001b[43m        \u001b[49m\u001b[43mhf_quantizer\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mhf_quantizer\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4029\u001b[0m \u001b[43m        \u001b[49m\u001b[43mkeep_in_fp32_modules\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeep_in_fp32_modules\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4030\u001b[0m \u001b[43m        \u001b[49m\u001b[43mgguf_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgguf_path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4031\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   4033\u001b[0m \u001b[38;5;66;03m# make sure token embedding weights are still tied if needed\u001b[39;00m\n\u001b[1;32m   4034\u001b[0m model\u001b[38;5;241m.\u001b[39mtie_weights()\n",
      "File \u001b[0;32m~/miniconda3/lib/python3.12/site-packages/transformers/modeling_utils.py:4502\u001b[0m, in \u001b[0;36mPreTrainedModel._load_pretrained_model\u001b[0;34m(cls, model, state_dict, loaded_keys, resolved_archive_file, pretrained_model_name_or_path, ignore_mismatched_sizes, sharded_metadata, _fast_init, low_cpu_mem_usage, device_map, offload_folder, offload_state_dict, dtype, hf_quantizer, keep_in_fp32_modules, gguf_path)\u001b[0m\n\u001b[1;32m   4498\u001b[0m                 set_module_tensor_to_device(\n\u001b[1;32m   4499\u001b[0m                     model_to_load, key, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcpu\u001b[39m\u001b[38;5;124m\"\u001b[39m, torch\u001b[38;5;241m.\u001b[39mempty(\u001b[38;5;241m*\u001b[39mparam\u001b[38;5;241m.\u001b[39msize(), dtype\u001b[38;5;241m=\u001b[39mdtype)\n\u001b[1;32m   4500\u001b[0m                 )\n\u001b[1;32m   4501\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 4502\u001b[0m         new_error_msgs, offload_index, state_dict_index \u001b[38;5;241m=\u001b[39m \u001b[43m_load_state_dict_into_meta_model\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   4503\u001b[0m \u001b[43m            \u001b[49m\u001b[43mmodel_to_load\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4504\u001b[0m \u001b[43m            \u001b[49m\u001b[43mstate_dict\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4505\u001b[0m \u001b[43m            \u001b[49m\u001b[43mstart_prefix\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4506\u001b[0m \u001b[43m            \u001b[49m\u001b[43mexpected_keys\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4507\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdevice_map\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice_map\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4508\u001b[0m \u001b[43m            \u001b[49m\u001b[43moffload_folder\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moffload_folder\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4509\u001b[0m \u001b[43m            \u001b[49m\u001b[43moffload_index\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moffload_index\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4510\u001b[0m \u001b[43m            \u001b[49m\u001b[43mstate_dict_folder\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstate_dict_folder\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4511\u001b[0m \u001b[43m            \u001b[49m\u001b[43mstate_dict_index\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstate_dict_index\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4512\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdtype\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4513\u001b[0m \u001b[43m            \u001b[49m\u001b[43mhf_quantizer\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mhf_quantizer\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4514\u001b[0m \u001b[43m            \u001b[49m\u001b[43mis_safetensors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mis_safetensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4515\u001b[0m \u001b[43m            \u001b[49m\u001b[43mkeep_in_fp32_modules\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeep_in_fp32_modules\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4516\u001b[0m \u001b[43m            \u001b[49m\u001b[43munexpected_keys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43munexpected_keys\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4517\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   4518\u001b[0m         error_msgs \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m new_error_msgs\n\u001b[1;32m   4519\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   4520\u001b[0m     \u001b[38;5;66;03m# Sharded checkpoint or whole but low_cpu_mem_usage==True\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/lib/python3.12/site-packages/transformers/modeling_utils.py:923\u001b[0m, in \u001b[0;36m_load_state_dict_into_meta_model\u001b[0;34m(model, state_dict, start_prefix, expected_keys, device_map, offload_folder, offload_index, state_dict_folder, state_dict_index, dtype, hf_quantizer, is_safetensors, keep_in_fp32_modules, unexpected_keys, pretrained_model_name_or_path)\u001b[0m\n\u001b[1;32m    921\u001b[0m             set_module_kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdtype\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mfloat32\n\u001b[1;32m    922\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 923\u001b[0m         param \u001b[38;5;241m=\u001b[39m \u001b[43mparam\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdtype\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    925\u001b[0m \u001b[38;5;66;03m# For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model, and which\u001b[39;00m\n\u001b[1;32m    926\u001b[0m \u001b[38;5;66;03m# uses `param.copy_(input_param)` that preserves the contiguity of the parameter in the model.\u001b[39;00m\n\u001b[1;32m    927\u001b[0m \u001b[38;5;66;03m# Reference: https://github.com/pytorch/pytorch/blob/db79ceb110f6646523019a59bbd7b838f43d4a86/torch/nn/modules/module.py#L2040C29-L2040C29\u001b[39;00m\n\u001b[1;32m    928\u001b[0m old_param \u001b[38;5;241m=\u001b[39m model\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "import contextlib\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "@contextlib.contextmanager\n",
    "def suppress_output():\n",
    "    \"\"\"Context manager to suppress stdout and stderr.\"\"\"\n",
    "    new_stdout, new_stderr = os.devnull, os.devnull\n",
    "    old_stdout, old_stderr = sys.stdout, sys.stderr\n",
    "    sys.stdout, sys.stderr = open(new_stdout, 'w'), open(new_stderr, 'w')\n",
    "    try:\n",
    "        yield\n",
    "    finally:\n",
    "        sys.stdout.close()\n",
    "        sys.stderr.close()\n",
    "        sys.stdout, sys.stderr = old_stdout, old_stderr\n",
    "\n",
    "\n",
    "def generate_proxy_tuning(model_base_name, model_tuned_name, model_target_name, input_text, max_length):\n",
    "    # Load the tokenizer once as we assume it's the same for all models\n",
    "    _, tokenizer = load_model_and_tokenizer(model_base_name)\n",
    "\n",
    "    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)\n",
    "    generated_tokens = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for _ in tqdm(range(max_length), desc=\"Generating tokens\"):\n",
    "            # Suppress output during model loading\n",
    "            with suppress_output():\n",
    "                model_base, _ = load_model_and_tokenizer(model_base_name)\n",
    "            logits_base = model_base(input_ids).logits\n",
    "            del model_base\n",
    "\n",
    "            with suppress_output():\n",
    "                model_tuned, _ = load_model_and_tokenizer(model_tuned_name)\n",
    "            logits_tuned = model_tuned(input_ids).logits\n",
    "            del model_tuned\n",
    "\n",
    "            with suppress_output():\n",
    "                model_target, _ = load_model_and_tokenizer(model_target_name)\n",
    "            logits_target = model_target(input_ids).logits\n",
    "            del model_target\n",
    "\n",
    "            # Calculate final logits\n",
    "            logits = logits_target + (logits_tuned - logits_base)\n",
    "\n",
    "            # Generate next token\n",
    "            predictions = torch.softmax(logits[:, -1, :], dim=-1)\n",
    "            next_token_id = torch.argmax(predictions).unsqueeze(0)\n",
    "            generated_tokens.append(next_token_id.item())\n",
    "\n",
    "            # Append the new token to the input sequence for the next iteration\n",
    "            input_ids = torch.cat([input_ids, next_token_id.unsqueeze(0)], dim=1)\n",
    "\n",
    "            # Break if end of sequence token is generated\n",
    "            if next_token_id.item() == tokenizer.eos_token_id:\n",
    "                break\n",
    "\n",
    "    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)\n",
    "    return generated_text\n",
    "\n",
    "\n",
    "generated = generate_proxy_tuning(\n",
    "    model_base_name=\"./autodl-tmp/Llama-2-7b-hf\",\n",
    "    model_tuned_name=\"./autodl-tmp/Llama-2-7b-chat-hf\",\n",
    "    model_target_name=\"./autodl-tmp/Llama-2-13b-hf\",\n",
    "    input_text=input_text,\n",
    "    max_length=256\n",
    ")\n",
    "\n",
    "print(\"\\033[5;34;46m\", generated, \"\\033[0m\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45084d26-a287-4b6b-ab38-ca7564a0f71c",
   "metadata": {},
   "source": [
    "正如我们所看到的，以前无法正确回答问题的 `Llama-2-13b-hf` 模型现在可以产生类似于通过 `Llama-2-13b-chat-hf` 创建的答案的答案。我们能够在不改变 `Llama-2-13b-hf` 模型的权重的情况下实现这一点！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e621ec31-97e7-48b0-9880-a69b31ef9937",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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
}