logits-processor-zoo_several_prompts

several_prompts.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e485c810-57a5-4514-9512-c720b0742f25",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:09:15.969307Z",
     "iopub.status.busy": "2025-04-25T14:09:15.968552Z",
     "iopub.status.idle": "2025-04-25T14:09:15.974401Z",
     "shell.execute_reply": "2025-04-25T14:09:15.973502Z",
     "shell.execute_reply.started": "2025-04-25T14:09:15.969259Z"
    }
   },
   "outputs": [],
   "source": [
    "# vLLM V1 does not currently accept logits processor so we need to disable it\n",
    "# https://docs.vllm.ai/en/latest/getting_started/v1_user_guide.html#deprecated-features\n",
    "import os\n",
    "\n",
    "os.environ[\"VLLM_USE_V1\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d81f2b4b-e460-45af-8450-d83730d95dfc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:09:15.976035Z",
     "iopub.status.busy": "2025-04-25T14:09:15.975506Z",
     "iopub.status.idle": "2025-04-25T14:09:20.284671Z",
     "shell.execute_reply": "2025-04-25T14:09:20.284186Z",
     "shell.execute_reply.started": "2025-04-25T14:09:15.975989Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-25 16:09:19 [__init__.py:239] Automatically detected platform cuda.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'0.8.4'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import vllm\n",
    "\n",
    "vllm.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "02c4ab8d-7eb9-44f1-956d-d0dba1ede372",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:09:20.285943Z",
     "iopub.status.busy": "2025-04-25T14:09:20.285583Z",
     "iopub.status.idle": "2025-04-25T14:09:20.288221Z",
     "shell.execute_reply": "2025-04-25T14:09:20.287837Z",
     "shell.execute_reply.started": "2025-04-25T14:09:20.285927Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1, 2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "135ce79c-aeba-4a19-8002-4c079ebda67d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:09:20.288822Z",
     "iopub.status.busy": "2025-04-25T14:09:20.288663Z",
     "iopub.status.idle": "2025-04-25T14:09:20.314503Z",
     "shell.execute_reply": "2025-04-25T14:09:20.313721Z",
     "shell.execute_reply.started": "2025-04-25T14:09:20.288809Z"
    }
   },
   "outputs": [],
   "source": [
    "model_id = \"Qwen/Qwen2.5-0.5B-Instruct\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "965ce9d5-4205-48e2-ad25-a1090b85cf3f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:09:20.316043Z",
     "iopub.status.busy": "2025-04-25T14:09:20.315531Z",
     "iopub.status.idle": "2025-04-25T14:09:20.336632Z",
     "shell.execute_reply": "2025-04-25T14:09:20.335525Z",
     "shell.execute_reply.started": "2025-04-25T14:09:20.316002Z"
    }
   },
   "outputs": [],
   "source": [
    "from logits_processor_zoo.vllm import ForceLastPhraseLogitsProcessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "68e1a236-5991-4221-b7b4-1a85532595d2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:09:20.338245Z",
     "iopub.status.busy": "2025-04-25T14:09:20.337722Z",
     "iopub.status.idle": "2025-04-25T14:10:04.144016Z",
     "shell.execute_reply": "2025-04-25T14:10:04.143264Z",
     "shell.execute_reply.started": "2025-04-25T14:09:20.338206Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-25 16:09:30 [config.py:689] This model supports multiple tasks: {'score', 'embed', 'generate', 'reward', 'classify'}. Defaulting to 'generate'.\n",
      "INFO 04-25 16:09:30 [config.py:1713] Defaulting to use mp for distributed inference\n",
      "INFO 04-25 16:09:30 [llm_engine.py:243] Initializing a V0 LLM engine (v0.8.4) with config: model='Qwen/Qwen2.5-0.5B-Instruct', speculative_config=None, tokenizer='Qwen/Qwen2.5-0.5B-Instruct', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=32768, download_dir='/big_storage/llms/hf_models/', load_format=LoadFormat.AUTO, tensor_parallel_size=2, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar', reasoning_backend=None), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=None, served_model_name=Qwen/Qwen2.5-0.5B-Instruct, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=None, chunked_prefill_enabled=False, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={\"splitting_ops\":[],\"compile_sizes\":[],\"cudagraph_capture_sizes\":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],\"max_capture_size\":256}, use_cached_outputs=False, \n",
      "WARNING 04-25 16:09:31 [multiproc_worker_utils.py:306] Reducing Torch parallelism from 18 threads to 1 to avoid unnecessary CPU contention. Set OMP_NUM_THREADS in the external environment to tune this value as needed.\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:31 [multiproc_worker_utils.py:225] Worker ready; awaiting tasks\n",
      "INFO 04-25 16:09:32 [cuda.py:292] Using Flash Attention backend.\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:32 [cuda.py:292] Using Flash Attention backend.\n",
      "INFO 04-25 16:09:33 [utils.py:993] Found nccl from library libnccl.so.2\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:33 [pynccl.py:69] vLLM is using nccl==2.21.5\n",
      "INFO 04-25 16:09:33 [utils.py:993] Found nccl from library libnccl.so.2\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:33 [pynccl.py:69] vLLM is using nccl==2.21.5\n",
      "INFO 04-25 16:09:33 [custom_all_reduce_utils.py:244] reading GPU P2P access cache from /home/asilva/.cache/vllm/gpu_p2p_access_cache_for_1, 2.json\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:33 [custom_all_reduce_utils.py:244] reading GPU P2P access cache from /home/asilva/.cache/vllm/gpu_p2p_access_cache_for_1, 2.json\n",
      "INFO 04-25 16:09:33 [shm_broadcast.py:264] vLLM message queue communication handle: Handle(local_reader_ranks=[1], buffer_handle=(1, 4194304, 6, 'psm_6ff2a2b3'), local_subscribe_addr='ipc:///tmp/9525e680-0572-4bec-b0fa-f2bea9a847a0', remote_subscribe_addr=None, remote_addr_ipv6=False)\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:33 [parallel_state.py:959] rank 1 in world size 2 is assigned as DP rank 0, PP rank 0, TP rank 1\n",
      "INFO 04-25 16:09:33 [parallel_state.py:959] rank 0 in world size 2 is assigned as DP rank 0, PP rank 0, TP rank 0\n",
      "INFO 04-25 16:09:33 [model_runner.py:1110] Starting to load model Qwen/Qwen2.5-0.5B-Instruct...\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:33 [model_runner.py:1110] Starting to load model Qwen/Qwen2.5-0.5B-Instruct...\n",
      "INFO 04-25 16:09:33 [weight_utils.py:265] Using model weights format ['*.safetensors']\n",
      "INFO 04-25 16:09:34 [weight_utils.py:315] No model.safetensors.index.json found in remote.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d47f70cdcae4c0f9925e3f207c06b2a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-25 16:09:34 [loader.py:458] Loading weights took 0.18 seconds\n",
      "INFO 04-25 16:09:34 [model_runner.py:1146] Model loading took 0.4643 GiB and 0.857723 seconds\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:34 [weight_utils.py:265] Using model weights format ['*.safetensors']\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:35 [weight_utils.py:315] No model.safetensors.index.json found in remote.\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:35 [loader.py:458] Loading weights took 0.16 seconds\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:35 [model_runner.py:1146] Model loading took 0.4643 GiB and 1.974101 seconds\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:38 [worker.py:267] Memory profiling takes 3.14 seconds\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:38 [worker.py:267] the current vLLM instance can use total_gpu_memory (15.73GiB) x gpu_memory_utilization (0.90) = 14.16GiB\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:38 [worker.py:267] model weights take 0.46GiB; non_torch_memory takes 0.17GiB; PyTorch activation peak memory takes 0.73GiB; the rest of the memory reserved for KV Cache is 12.79GiB.\n",
      "INFO 04-25 16:09:39 [worker.py:267] Memory profiling takes 3.24 seconds\n",
      "INFO 04-25 16:09:39 [worker.py:267] the current vLLM instance can use total_gpu_memory (15.73GiB) x gpu_memory_utilization (0.90) = 14.16GiB\n",
      "INFO 04-25 16:09:39 [worker.py:267] model weights take 0.46GiB; non_torch_memory takes 0.17GiB; PyTorch activation peak memory takes 1.44GiB; the rest of the memory reserved for KV Cache is 12.08GiB.\n",
      "INFO 04-25 16:09:39 [executor_base.py:112] # cuda blocks: 131958, # CPU blocks: 43690\n",
      "INFO 04-25 16:09:39 [executor_base.py:117] Maximum concurrency for 32768 tokens per request: 64.43x\n",
      "INFO 04-25 16:09:41 [model_runner.py:1456] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:09:41 [model_runner.py:1456] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "806e4fbfaa02461296ce93c17724d793",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Capturing CUDA graph shapes:   0%|          | 0/35 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:10:01 [custom_all_reduce.py:195] Registering 1715 cuda graph addresses\n",
      "INFO 04-25 16:10:01 [custom_all_reduce.py:195] Registering 1715 cuda graph addresses\n",
      "INFO 04-25 16:10:01 [model_runner.py:1598] Graph capturing finished in 20 secs, took 0.75 GiB\n",
      "\u001b[1;36m(VllmWorkerProcess pid=3626680)\u001b[0;0m INFO 04-25 16:10:01 [model_runner.py:1598] Graph capturing finished in 20 secs, took 0.75 GiB\n",
      "INFO 04-25 16:10:01 [llm_engine.py:449] init engine (profile, create kv cache, warmup model) took 26.23 seconds\n"
     ]
    }
   ],
   "source": [
    "model = vllm.LLM(model=model_id, download_dir=\"/big_storage/llms/hf_models/\", tensor_parallel_size=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "11b1991d-d954-4707-944c-e34b5fb87abb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:10:04.145940Z",
     "iopub.status.busy": "2025-04-25T14:10:04.145434Z",
     "iopub.status.idle": "2025-04-25T14:10:04.150503Z",
     "shell.execute_reply": "2025-04-25T14:10:04.149794Z",
     "shell.execute_reply.started": "2025-04-25T14:10:04.145893Z"
    }
   },
   "outputs": [],
   "source": [
    "tokenizer = model.get_tokenizer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2a40eba8-75f2-4cb8-bde6-33e6c3959ab4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:10:04.152416Z",
     "iopub.status.busy": "2025-04-25T14:10:04.151655Z",
     "iopub.status.idle": "2025-04-25T14:10:04.181109Z",
     "shell.execute_reply": "2025-04-25T14:10:04.180012Z",
     "shell.execute_reply.started": "2025-04-25T14:10:04.152370Z"
    }
   },
   "outputs": [],
   "source": [
    "logits_processors = [\n",
    "    ForceLastPhraseLogitsProcessor(\"\\n\\nWait\", tokenizer)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "90496c07-5722-48d2-9dae-536a68f12364",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:10:04.182879Z",
     "iopub.status.busy": "2025-04-25T14:10:04.182422Z",
     "iopub.status.idle": "2025-04-25T14:10:04.195267Z",
     "shell.execute_reply": "2025-04-25T14:10:04.194192Z",
     "shell.execute_reply.started": "2025-04-25T14:10:04.182835Z"
    }
   },
   "outputs": [],
   "source": [
    "user_prompts = [\n",
    "    \"Hello!\",\n",
    "    \"What's the capital of France?\",\n",
    "    \"Who was the first president of the United States?\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "945d8154-9434-4ab3-b4fc-c740a2e25d98",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:10:04.196381Z",
     "iopub.status.busy": "2025-04-25T14:10:04.196165Z",
     "iopub.status.idle": "2025-04-25T14:10:04.213620Z",
     "shell.execute_reply": "2025-04-25T14:10:04.212631Z",
     "shell.execute_reply.started": "2025-04-25T14:10:04.196361Z"
    }
   },
   "outputs": [],
   "source": [
    "def create_prompt(prompt):\n",
    "    messages = [{\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},{\"role\": \"user\", \"content\": prompt}]\n",
    "    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bf5019e7-f7a5-435d-ae20-09a3a973d7b9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:10:04.215250Z",
     "iopub.status.busy": "2025-04-25T14:10:04.214797Z",
     "iopub.status.idle": "2025-04-25T14:10:04.265792Z",
     "shell.execute_reply": "2025-04-25T14:10:04.264660Z",
     "shell.execute_reply.started": "2025-04-25T14:10:04.215208Z"
    }
   },
   "outputs": [],
   "source": [
    "prompts = [create_prompt(prompt) for prompt in user_prompts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ec2500d2-4e5e-4fa2-8369-fa2811d3663c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:10:04.269619Z",
     "iopub.status.busy": "2025-04-25T14:10:04.269147Z",
     "iopub.status.idle": "2025-04-25T14:10:04.879972Z",
     "shell.execute_reply": "2025-04-25T14:10:04.879335Z",
     "shell.execute_reply.started": "2025-04-25T14:10:04.269583Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5c30c30134ce480d96678643345951ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processed prompts:   0%|          | 0/3 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gen_output = model.generate(\n",
    "            prompts,\n",
    "            vllm.SamplingParams(\n",
    "                n=1,\n",
    "                temperature=0,\n",
    "                seed=0,\n",
    "                skip_special_tokens=True,\n",
    "                max_tokens=200,\n",
    "                logits_processors=logits_processors\n",
    "            ),\n",
    "            use_tqdm=True\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "91600c83-5544-4623-8697-c06d195d6c6f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-25T14:10:04.880721Z",
     "iopub.status.busy": "2025-04-25T14:10:04.880553Z",
     "iopub.status.idle": "2025-04-25T14:10:04.884768Z",
     "shell.execute_reply": "2025-04-25T14:10:04.884329Z",
     "shell.execute_reply.started": "2025-04-25T14:10:04.880705Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: 'Hello! How can I help you today?'\n",
      "Generated text: 'The capital of France is Paris.\\n\\nParis is the largest city in France and the second-largest city in the European Union. It is located on the left bank of the Seine River, which flows through the city. The city is known for its rich history, art, and culture, and is home to many famous landmarks such as the Eiffel Tower, Louvre Museum, and Notre-Dame Cathedral.'\n",
      "Generated text: 'The first president of the United StatesWaiter was George Washington.'\n"
     ]
    }
   ],
   "source": [
    "for output in gen_output:\n",
    "    prompt = output.prompt\n",
    "    generated_text = output.outputs[0].text\n",
    "    print(f\"Generated text: {generated_text!r}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5011cf2-f26e-4a00-8bcd-541b006d5cd3",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}