322 lines
8.3 KiB
Plaintext
322 lines
8.3 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
|
|
"Unsloth: Failed to patch Gemma3ForConditionalGeneration.\n",
|
|
"🦥 Unsloth Zoo will now patch everything to make training faster!\n",
|
|
"INFO 04-23 16:42:57 [__init__.py:239] Automatically detected platform cuda.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Unsloth should be imported before transformers to ensure all optimizations are applied.\n",
|
|
"from unsloth import FastLanguageModel"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"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": 3,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"temperature = 0.6\n",
|
|
"model_path = \"/root/models/notebook_merged_4bit\"\n",
|
|
"xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Load tokenizer\n",
|
|
"tokenizer = AutoTokenizer.from_pretrained(\n",
|
|
" model_path,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"==((====))== Unsloth 2025.3.19: Fast Qwen2 patching. Transformers: 4.51.3. vLLM: 0.8.2.\n",
|
|
" \\\\ /| NVIDIA RTX 6000 Ada Generation. Num GPUs = 1. Max memory: 47.5 GB. Platform: Linux.\n",
|
|
"O^O/ \\_/ \\ Torch: 2.6.0+cu124. CUDA: 8.9. CUDA Toolkit: 12.4. Triton: 3.2.0\n",
|
|
"\\ / Bfloat16 = TRUE. FA [Xformers = 0.0.29.post2. FA2 = False]\n",
|
|
" \"-____-\" Free license: http://github.com/unslothai/unsloth\n",
|
|
"Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Sliding Window Attention is enabled but not implemented for `eager`; unexpected results may be encountered.\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"validate_env\n",
|
|
"device_map sequential\n",
|
|
"validate_env\n",
|
|
"device_map OrderedDict([('', 0)])\n"
|
|
]
|
|
},
|
|
{
|
|
"data": {
|
|
"application/vnd.jupyter.widget-view+json": {
|
|
"model_id": "4f302eb9995d47fab2aa6339aa00a8d8",
|
|
"version_major": 2,
|
|
"version_minor": 0
|
|
},
|
|
"text/plain": [
|
|
"Loading checkpoint shards: 0%| | 0/4 [00:00<?, ?it/s]"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"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",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"Qwen2ForCausalLM(\n",
|
|
" (model): Qwen2Model(\n",
|
|
" (embed_tokens): Embedding(152064, 5120, padding_idx=151654)\n",
|
|
" (layers): ModuleList(\n",
|
|
" (0-63): 64 x Qwen2DecoderLayer(\n",
|
|
" (self_attn): Qwen2Attention(\n",
|
|
" (q_proj): Linear4bit(in_features=5120, out_features=5120, bias=True)\n",
|
|
" (k_proj): Linear4bit(in_features=5120, out_features=1024, bias=True)\n",
|
|
" (v_proj): Linear4bit(in_features=5120, out_features=1024, bias=True)\n",
|
|
" (o_proj): Linear4bit(in_features=5120, out_features=5120, bias=False)\n",
|
|
" (rotary_emb): LlamaRotaryEmbedding()\n",
|
|
" )\n",
|
|
" (mlp): Qwen2MLP(\n",
|
|
" (gate_proj): Linear4bit(in_features=5120, out_features=27648, bias=False)\n",
|
|
" (up_proj): Linear4bit(in_features=5120, out_features=27648, bias=False)\n",
|
|
" (down_proj): Linear4bit(in_features=27648, out_features=5120, bias=False)\n",
|
|
" (act_fn): SiLU()\n",
|
|
" )\n",
|
|
" (input_layernorm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
|
|
" (post_attention_layernorm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
|
|
" )\n",
|
|
" )\n",
|
|
" (norm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
|
|
" (rotary_emb): LlamaRotaryEmbedding()\n",
|
|
" )\n",
|
|
" (lm_head): Linear(in_features=5120, out_features=152064, bias=False)\n",
|
|
")"
|
|
]
|
|
},
|
|
"execution_count": 6,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"# enable unsloth optimizations\n",
|
|
"FastLanguageModel.for_inference(model)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Device set to use cuda:0\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# 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",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 8,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 9,
|
|
"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": 10,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"text = tokenizer.apply_chat_template(\n",
|
|
" messages,\n",
|
|
" tokenize=False,\n",
|
|
" add_generation_prompt=False,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 11,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"streamer = TextIteratorStreamer(\n",
|
|
" tokenizer,\n",
|
|
" skip_prompt=True,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 12,
|
|
"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": 13,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"generation_kwargs[\"logits_processor\"] = [logits_processor.copy()]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 14,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"generation_thread = Thread(\n",
|
|
" target=pipeline,\n",
|
|
" kwargs=generation_kwargs\n",
|
|
")\n",
|
|
"\n",
|
|
"generation_thread.start()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 15,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"<write_stdout>Hi, I'm an AI assistant. I don't have feelings, but I'm here to help you. How can I assist you today?</write_stdout><|im_end|>"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"for text in streamer:\n",
|
|
" print(text, end=\"\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 16,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"\n",
|
|
"generation_thread.join()"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "sia",
|
|
"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.10.12"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|