diff --git a/sia/llm_engine/qwq_vllm.ipynb b/sia/llm_engine/qwq_vllm.ipynb new file mode 100644 index 0000000..87c9621 --- /dev/null +++ b/sia/llm_engine/qwq_vllm.ipynb @@ -0,0 +1,255 @@ +{ + "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 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\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 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", + "\n", + "\n", + "\n", + "Hello! I'm just a computer program, but I'm here to help you. How can I assist you today?\n", + "" + ] + } + ], + "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 +}