llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

New web interface, move llm engine to separate process

Browse Source
This commit is contained in:
geens
2025-05-20 09:43:17 +02:00
parent 895a533e01
commit d4a4902b94
137 changed files with 4850 additions and 3503 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
tools/old_llm_engines/hf_llm_engine.py Normal file
View File

@@ -0,0 +1,87 @@
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoConfig
from typing import Iterator, Optional, Callable
from . import LlmEngine
class HfLlmEngine(LlmEngine):
"""
LLM Engine implementation using HuggingFace's InferenceClient.
"""
def __init__(
self,
model: str,
temperature: float,
api_token: Optional[str],
):
"""
Initialize the HuggingFace Inference API LLM Engine.
Args:
model: HuggingFace model ID to use
temperature: Sampling temperature
api_token: HuggingFace API token
"""
self._model = model
self._temperature = temperature
self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
self._config = AutoConfig.from_pretrained(model, token=api_token)
self._client = InferenceClient(token=api_token)
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
"""
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
continuation_text: Part of the response that is already generated
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context},
{"role": "assistant", "content": continuation_text},
]
stream = self._client.chat_completion(
model=self._model,
messages=messages,
temperature=self._temperature,
add_generation_prompt=False,
stream=True
)
try:
for response in stream:
if should_stop():
stream.close()
break
if content := response.choices[0].delta.content:
yield content
finally:
stream.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(prompt))
def token_limit(self) -> int:
"""
Get the model's context window size.
Returns:
int: Maximum number of tokens the model can process
"""
return self._config.max_position_embeddings

137
tools/old_llm_engines/local_llm_engine.py Normal file
View File

@@ -0,0 +1,137 @@
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
from typing import Iterator, Optional, Callable
from xml_schema_validator import XmlLogitsProcessor
import sys
import torch
from . import LlmEngine
from .. import util
class LocalLlmEngine(LlmEngine):
def __init__(
self,
model_path: str,
temperature: float,
token_limit: int,
xml_schema_text: Optional[str] = None,
api_token: Optional[str] = None,
):
"""
Initialize the LLM Engine with a model path.
Args:
model_path: Path to the model weights to be used.
temperature: Temperature for sampling
token_limit: Maximum number of tokens to generate
xml_schema_text: Optional XML schema to validate against
api_token: Huggingface API key
"""
self._temperature = temperature
self._token_limit = token_limit
self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
model = AutoModelForCausalLM.from_pretrained(
model_path,
return_dict=True,
low_cpu_mem_usage=True,
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True,
token=api_token,
)
if self._tokenizer.pad_token_id is None:
self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
if model.config.pad_token_id is None:
model.config.pad_token_id = model.config.eos_token_id
self._pipeline = pipeline(
"text-generation",
model=model,
tokenizer=self._tokenizer,
torch_dtype=torch.bfloat16,
device_map="auto",
return_full_text=False,
)
if xml_schema_text:
self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
else:
self._logits_processor = None
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
"""
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
continuation_text: Part of the response that is already generated
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context},
{"role": "assistant", "content": continuation_text},
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False,
add_generation_prompt=False,
)
streamer = TextIteratorStreamer(
self._tokenizer,
skip_prompt=True
)
generation_kwargs = {
"text_inputs": prompt,
"do_sample": True,
"temperature": self._temperature,
"max_new_tokens": self.token_limit(),
"streamer": streamer,
}
if self._logits_processor:
generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
generation_thread = Thread(
target=self._pipeline,
kwargs=generation_kwargs
)
generation_thread.start()
for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
yield text
if should_stop():
break
generation_thread.join()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""
Count tokens for the given system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string
Returns:
int: Total number of tokens
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(prompt))
def token_limit(self) -> int:
"""
Get the model's context window size.
Returns:
int: Maximum number of tokens the model can process
"""
if self._token_limit is not None:
return self._token_limit
else:
return self._pipeline.model.config.max_position_embeddings

77
tools/old_llm_engines/openai_llm_engine.py Normal file
View File

@@ -0,0 +1,77 @@
from typing import Callable, Iterator
import openai
import tiktoken
from . import LlmEngine
class OpenAILlmEngine(LlmEngine):
"""
LLM Engine implementation using OpenAI's API.
Supports streaming responses from chat completion models.
"""
def __init__(
self,
model: str,
temperature: float,
token_limit: int,
api_key: str,
):
"""
Initialize the OpenAI LLM Engine.
Args:
model: OpenAI model to use
temperature: Temperature for sampling
api_key: OpenAI API key
token_limit: Maximum number of tokens to generate
"""
self._model = model
self._temperature = temperature
self._token_limit = token_limit
self._client = openai.Client(
api_key=api_key,
)
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
if continuation_text:
print("OpenAI LLM Engine: continuation_text is not supported")
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
stream = self._client.chat.completions.create(
model=self._model,
messages=messages,
temperature=self._temperature,
stream=True,
)
try:
for chunk in stream:
if should_stop():
break
if content := chunk.choices[0].delta.content:
yield content
finally:
stream.close()
#stream.response.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""
Calculate the total token count for the system prompt and context.
Args:
system_prompt: The system prompt string
main_context: The main context string
Returns:
int: Total number of tokens
"""
encoding = tiktoken.encoding_for_model(self._model)
return len(encoding.encode(system_prompt)) + len(encoding.encode(main_context))
def token_limit(self) -> int:
return self._token_limit

227
tools/old_llm_engines/qwq_llm_engine.ipynb Normal file
View File

@@ -0,0 +1,227 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"source /root/venvs/sia/bin/activate\n",
"apt-get update && apt-get install -y cuda-toolkit\n",
"pip install flash-attn --no-build-isolation"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Unsloth should be imported before transformers to ensure all optimizations are applied.\n",
"from unsloth import FastLanguageModel"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"from threading import Thread\n",
"from transformers import AutoTokenizer, TextIteratorStreamer, pipeline\n",
"from xml_schema_validator import XmlLogitsProcessor"
]
},
{
"cell_type": "code",
"execution_count": null,
"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": null,
"metadata": {},
"outputs": [],
"source": [
"# Load tokenizer\n",
"tokenizer = AutoTokenizer.from_pretrained(\n",
" model_path,\n",
" legacy=False,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Load model\n",
"model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
" model_path,\n",
" gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
" load_in_4bit=True,\n",
" attn_implementation=\"flash_attention_2\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# enable unsloth optimizations\n",
"FastLanguageModel.for_inference(model)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"# Create inference pipeline with memory-efficient settings\n",
"pipeline = pipeline(\n",
" \"text-generation\",\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" return_full_text=False,\n",
" torch_dtype=torch.float16,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)"
]
},
{
"cell_type": "code",
"execution_count": null,
"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": null,
"metadata": {},
"outputs": [],
"source": [
"text = tokenizer.apply_chat_template(\n",
" messages,\n",
" tokenize=False,\n",
" add_generation_prompt=False,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"streamer = TextIteratorStreamer(\n",
" tokenizer,\n",
" skip_prompt=True,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generation_kwargs = {\n",
" \"text_inputs\": text,\n",
" \"do_sample\": True,\n",
" \"temperature\": temperature,\n",
" \"streamer\": streamer,\n",
" \"use_cache\": True,\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generation_kwargs[\"logits_processor\"] = [logits_processor.copy()]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generation_thread = Thread(\n",
" target=pipeline,\n",
" kwargs=generation_kwargs\n",
")\n",
"\n",
"generation_thread.start()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for text in streamer:\n",
" print(text, end=\"\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"generation_thread.join()"
]
}
],
"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
}

135
tools/old_llm_engines/qwq_llm_engine.py Normal file
View File

@@ -0,0 +1,135 @@
# Unsloth should be imported before transformers to ensure all optimizations are applied.
from unsloth import FastLanguageModel
from pathlib import Path
from threading import Thread
from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
from typing import Callable, Iterator, Optional
from xml_schema_validator import XmlLogitsProcessor
from . import LlmEngine
from .. import util
class QwQLlmEngine(LlmEngine):
def __init__(
self,
model_path: Path,
temperature: float,
xml_schema_text: Optional[str] = None,
):
"""
Initialize the QwQ LLM Engine.
Args:
model_path: Local path to the model
temperature: Sampling temperature
xml_schema_text: Optional XML schema to validate against
"""
self._temperature = temperature
# Load tokenizer
self._tokenizer = AutoTokenizer.from_pretrained(
model_path,
)
# Load model
self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
model_path,
gpu_memory_utilization = 0.5, # Reduce if out of memory
)
# enable unsloth optimizations
FastLanguageModel.for_inference(self._model)
# Create inference pipeline
self._pipeline = pipeline(
"text-generation",
model=self._model,
tokenizer=self._tokenizer,
return_full_text=False,
)
if xml_schema_text:
self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
else:
self._logits_processor = None
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
"""
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
continuation_text: Part of the response that is already generated
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context},
{"role": "assistant", "content": continuation_text},
]
text = self._tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=False,
)
streamer = TextIteratorStreamer(
self._tokenizer,
skip_prompt=True,
)
generation_kwargs = {
"text_inputs": text,
"do_sample": True,
"temperature": self._temperature,
"streamer": streamer,
"use_cache": True,
}
if self._logits_processor:
generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
generation_thread = Thread(
target=self._pipeline,
kwargs=generation_kwargs
)
generation_thread.start()
for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
yield text
if should_stop():
break
generation_thread.join()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""
Count tokens for the given system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string
Returns:
int: Total number of tokens
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(prompt))
def token_limit(self) -> int:
return self._pipeline.model.config.max_position_embeddings

255
tools/old_llm_engines/qwq_vllm.ipynb Normal file
View File

@@ -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 <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
}