fix a cli bug

Author: harveyp123

medusa/inference/cli.py

@@ -37,6 +37,7 @@ def main(args):
         model = MedusaModel.from_pretrained(
             args.model,
             args.base_model,
+            medusa_num_heads = 4,
             torch_dtype=torch.float16,
             low_cpu_mem_usage=True,
             device_map="auto",

medusa/model/medusa_model.py

@@ -126,8 +126,8 @@ def get_tokenizer(self):
     def from_pretrained(
         cls,
         medusa_head_name_or_path,
-        medusa_num_heads=None,
         base_model=None,
+        medusa_num_heads=None,
         **kwargs,
     ):
         """

notebooks/medusa_configuration_explained.ipynb

@@ -121,7 +121,7 @@
     {
      "data": {
       "text/plain": [
-       "<matplotlib.image.AxesImage at 0x7f79c4e723d0>"
+       "<matplotlib.image.AxesImage at 0x7f1b60e80520>"
       ]
      },
      "execution_count": 6,

notebooks/medusa_inference_explained.ipynb

@@ -11,7 +11,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -42,7 +42,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -151,9 +151,38 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Overriding medusa_num_heads as: 4\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ef69040c760f4e4b949e27b2c09526d2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama.LlamaTokenizer'>. If you see this, DO NOT PANIC! This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thouroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n"
+     ]
+    }
+   ],
    "source": [
     "model_name = 'FasterDecoding/medusa-vicuna-7b-v1.3'\n",
     "model = MedusaModel.from_pretrained(\n",
@@ -180,7 +209,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -200,7 +229,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -218,9 +247,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 9,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Output length: 403\n",
+      "Compression ratio: tensor(2.4724, device='cuda:0')\n"
+     ]
+    }
+   ],
    "source": [
     "with torch.inference_mode():\n",
     "    input_ids = tokenizer([prompt]).input_ids\n",
@@ -249,9 +287,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Once upon a time, in a small village nestled in the Andes mountains, there lived a charming llama named Luna. Luna was known for her kind heart and her love of coffee. She would often spend her afternoons sipping on a steaming cup of joe at the local café, chatting with the villagers and enjoying the warmth of the sun on her back.\n",
+      "\n",
+      "One day, as Luna was grazing on some fresh grass, she noticed that her hair was starting to grow longer and thicker. At first, she didn't think much of it, but as the days went on, her hair continued to grow and change. It became thick and wiry, with sharp spikes protruding from it.\n",
+      "\n",
+      "Luna was confused and a little scared by her new appearance. She had always been a gentle creature, and now she looked like a monster. She knew that she couldn't stay in the village anymore, so she set off on a journey to find a new home.\n",
+      "\n",
+      "As she wandered through the mountains, Luna stumbled upon a beautiful clearing. In the center of the clearing stood a small cottage, with a sign hanging outside that read \"Café Llama.\" Luna knew that this was where she belonged.\n",
+      "\n",
+      "She transformed the cottage into a cozy coffee shop, serving the best coffee in the mountains. The villagers were amazed by Luna's transformation, and they flocked to her café to taste her delicious brews.\n",
+      "\n",
+      "Luna's Medusa-like hair became her signature style, and she quickly became known as the most charming llama in the land. She spent her days sipping coffee, chatting with customers, and enjoying the warmth of the sun on her back. And she knew that she had finally found her true home.</s>\n"
+     ]
+    }
+   ],
    "source": [
     "output = tokenizer.decode(\n",
     "                    output_ids,\n",
@@ -275,9 +331,30 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "==================================================\n",
+      "Wall time init:                              0.026\n",
+      "Wall time medusa:                            0.031\n",
+      "Wall time Tree:                              3.732\n",
+      "Wall time Posterior:                         0.025\n",
+      "Wall time Update:                            0.051\n",
+      "--------------------------------------------------\n",
+      "Wall time portion medusa:                    0.008\n",
+      "Wall time portion Tree:                      0.965\n",
+      "Wall time portion Posterior:                 0.007\n",
+      "Wall time portion Update:                    0.013\n",
+      "--------------------------------------------------\n",
+      "Tokens/second:                             104.247\n",
+      "==================================================\n"
+     ]
+    }
+   ],
    "source": [
     "max_length = 50\n",
     "\n",
@@ -307,13 +384,6 @@
     "print(format_string(\"Tokens/second: \", new_token / time_total, max_length))\n",
     "print('='*max_length)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

notebooks/medusa_introduction.ipynb

@@ -66,10 +66,17 @@
    "execution_count": 2,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Overriding medusa_num_heads as: 4\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b6655fe35f99442e9c250f18882ce883",
+       "model_id": "f4d7c1aaf692402f959920bc9ccd8593",
        "version_major": 2,
        "version_minor": 0
       },
@@ -234,14 +241,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "['a']\n"
+      "['a']\n",
+      "['time']\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "['time']\n",
       "[',']\n",
       "['in']\n",
       "['a']\n",
@@ -917,7 +924,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Prediction @ 1: ['Once']\n",
+      "Prediction @ 1: ['Once']\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "Prediction @ 2: ['upon', 'a']\n",
       "Prediction @ 3: ['time', ',', 'in']\n",
       "Prediction @ 4: ['a', 'small']\n",
@@ -1583,7 +1596,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 29,
    "metadata": {},
    "outputs": [
     {