{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# vLLM Streaming Implementation\n",
    "\n",
    "This notebook demonstrates how to implement streaming capability with vLLM, comparable to the unsloth implementation.\n",
    "\n",
    "First, let's make sure we have vLLM installed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-25 19:36:31 [__init__.py:239] Automatically detected platform cuda.\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "from vllm import SamplingParams\n",
    "from transformers import AutoTokenizer\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "temperature = 0.6\n",
    "model_path = \"/root/models/current\"\n",
    "xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_path,\n",
    "    legacy=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"Always respond in a <write_stdout></write_stdout> xml block.\"},\n",
    "    {\"role\": \"user\", \"content\": \"Hi, how are you?\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define sampling parameters\n",
    "sampling_params = SamplingParams(\n",
    "    temperature=temperature,\n",
    "    top_p=0.95,\n",
    "    max_tokens=512,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-25 19:36:40 [config.py:585] This model supports multiple tasks: {'generate', 'score', 'reward', 'embed', 'classify'}. Defaulting to 'generate'.\n",
      "WARNING 04-25 19:36:42 [config.py:664] bitsandbytes quantization is not fully optimized yet. The speed can be slower than non-quantized models.\n",
      "WARNING 04-25 19:36:42 [arg_utils.py:1854] --quantization bitsandbytes is not supported by the V1 Engine. Falling back to V0. \n",
      "INFO 04-25 19:36:42 [llm_engine.py:241] Initializing a V0 LLM engine (v0.8.2) with config: model='/root/models/current', speculative_config=None, tokenizer='/root/models/current', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=4096, download_dir=None, load_format=bitsandbytes, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=bitsandbytes, 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=/root/models/current, 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",
      "INFO 04-25 19:36:42 [cuda.py:291] Using Flash Attention backend.\n",
      "INFO 04-25 19:36:43 [parallel_state.py:954] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0\n",
      "INFO 04-25 19:36:43 [model_runner.py:1110] Starting to load model /root/models/current...\n",
      "INFO 04-25 19:36:43 [loader.py:1155] Loading weights with BitsAndBytes quantization. May take a while ...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8b9f3cb293484cac932e6cedd841c813",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54f8aa5eefdb43d8bc07274044a8bc1c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-25 19:36:51 [model_runner.py:1146] Model loading took 18.0523 GB and 8.113452 seconds\n",
      "INFO 04-25 19:36:55 [worker.py:267] Memory profiling takes 3.23 seconds\n",
      "INFO 04-25 19:36:55 [worker.py:267] the current vLLM instance can use total_gpu_memory (47.53GiB) x gpu_memory_utilization (0.90) = 42.78GiB\n",
      "INFO 04-25 19:36:55 [worker.py:267] model weights take 18.05GiB; non_torch_memory takes 0.06GiB; PyTorch activation peak memory takes 1.59GiB; the rest of the memory reserved for KV Cache is 23.08GiB.\n",
      "INFO 04-25 19:36:55 [executor_base.py:111] # cuda blocks: 5907, # CPU blocks: 1024\n",
      "INFO 04-25 19:36:55 [executor_base.py:116] Maximum concurrency for 4096 tokens per request: 23.07x\n",
      "INFO 04-25 19:36:58 [model_runner.py:1442] 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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Capturing CUDA graph shapes: 100%|██████████| 35/35 [01:01<00:00,  1.75s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-25 19:38:00 [model_runner.py:1570] Graph capturing finished in 61 secs, took 1.98 GiB\n",
      "INFO 04-25 19:38:00 [llm_engine.py:447] init engine (profile, create kv cache, warmup model) took 68.57 seconds\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from vllm import LLM, SamplingParams\n",
    "import time\n",
    "\n",
    "# Initialize LLM\n",
    "llm = LLM(\n",
    "    model=model_path,\n",
    "    tensor_parallel_size=1,\n",
    "    max_model_len=4096,\n",
    "    quantization=\"bitsandbytes\",\n",
    "    load_format=\"bitsandbytes\",\n",
    "    trust_remote_code=True,\n",
    "    # Enable streaming\n",
    "    enable_lora=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting generation with token-by-token output:\n",
      "<think>\n",
      "Okay, the user greeted me with \"Hi, how are you?\" I need to respond appropriately. Let me see... The instructions say to always use the <write_stdout> XML tag. So first, I should acknowledge their greeting and state that I'm an AI, then ask how I can assist them. Keep it friendly and helpful. Let me make sure I don't add any extra information beyond that. Just a simple response. Alright, that should work.\n",
      "</think>\n",
      "\n",
      "<write_stdout>\n",
      "Hello! I'm just a computer program, but I'm here to help you. How can I assist you today?\n",
      "</write_stdout>"
     ]
    }
   ],
   "source": [
    "previous_text = \"\"\n",
    "print(\"Starting generation with token-by-token output:\")\n",
    "\n",
    "# Try with direct iteration over the generator\n",
    "for output in llm.generate(prompt, sampling_params, use_tqdm=False):\n",
    "    if hasattr(output, 'outputs') and output.outputs and len(output.outputs) > 0:\n",
    "        generated_text = output.outputs[0].text\n",
    "        if len(generated_text) > len(previous_text):\n",
    "            new_text = generated_text[len(previous_text):]\n",
    "            sys.stdout.write(new_text)\n",
    "            sys.stdout.flush()\n",
    "            previous_text = generated_text"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sia",
   "language": "python",
   "name": "sia"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}