Fixed inference notebook

tools/train/train/qwq_infer.ipynb

@@ -2,25 +2,19 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
-     "languageId": "plaintext"
-    }
-   },
    "outputs": [],
    "source": [
-    "from transformers import AutoModelForCausalLM, AutoTokenizer"
+    "from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig\n",
+    "from threading import Thread\n",
+    "import torch"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
-     "languageId": "plaintext"
-    }
-   },
    "outputs": [],
    "source": [
     "model_path = \"/root/models/notebook\""
@@ -28,30 +22,60 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
-     "languageId": "plaintext"
-    }
-   },
    "outputs": [],
    "source": [
-    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "quantization_config = BitsAndBytesConfig(\n",
-    "    model_path,\n",
+    "    load_in_4bit=True,\n",
-    "    return_dict=True,\n",
+    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
-    "    torch_dtype=\"auto\",\n",
+    "    bnb_4bit_quant_type=\"nf4\",\n",
-    "    device_map=\"auto\"\n",
+    "    bnb_4bit_use_double_quant=True\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
+   "outputs": [
-     "languageId": "plaintext"
+    {
-    }
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5f35aa01aafa4bc2beb19f437373945e",
+       "version_major": 2,
+       "version_minor": 0
       },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/14 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    model_path,\n",
+    "    return_dict=True,\n",
+    "    device_map=\"auto\",\n",
+    "    attn_implementation=\"flash_attention_2\",\n",
+    "    use_cache=True,\n",
+    "    quantization_config=quantization_config,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "pip install flash-attn --no-build-isolation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
    "outputs": [],
    "source": [
     "tokenizer = AutoTokenizer.from_pretrained(model_path)"
@@ -59,13 +83,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
+   "outputs": [
-     "languageId": "plaintext"
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Device set to use cuda:0\n"
+     ]
     }
-   },
+   ],
-   "outputs": [],
    "source": [
     "pipeline = pipeline(\n",
     "    \"text-generation\",\n",
@@ -77,12 +105,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
-     "languageId": "plaintext"
-    }
-   },
    "outputs": [],
    "source": [
     "prompt = \"How many r's are in the word \\\"strawberry\\\"\"\n",
@@ -93,12 +117,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
-     "languageId": "plaintext"
-    }
-   },
    "outputs": [],
    "source": [
     "text = tokenizer.apply_chat_template(\n",
@@ -110,12 +130,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 9,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
-     "languageId": "plaintext"
-    }
-   },
    "outputs": [],
    "source": [
     "streamer = TextIteratorStreamer(\n",
@@ -126,19 +142,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
-     "languageId": "plaintext"
-    }
-   },
    "outputs": [],
    "source": [
-    "generation_thread = Thread(target=self._pipeline, kwargs=dict(\n",
+    "generation_thread = Thread(target=pipeline, kwargs=dict(\n",
     "    text_inputs=text,\n",
     "    do_sample=True,\n",
     "    temperature=0.7,\n",
-    "    max_new_tokens=512,\n",
+    "    max_new_tokens=2048,\n",
     "    streamer=streamer,\n",
     "))\n",
     "generation_thread.start()"
@@ -146,23 +158,123 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
-   "metadata": {
+   "metadata": {},
-    "vscode": {
+   "outputs": [
-     "languageId": "plaintext"
+    {
-    }
+     "name": "stdout",
-   },
+     "output_type": "stream",
-   "outputs": [],
+     "text": [
-   "source": [
+      "Okay, let's see. The user is asking how many 'r's are in the word \"strawberry.\" Hmm, I need to make sure I get this right. First, I should probably write out the word letter by letter to check each letter. Let me spell it out: S-T-R-A-W-B-E-R-R-Y. Wait, let me go slow. S, T, R, A, W, B, E, R, R, Y. Let me count the 'R's again.\n",
       "\n",
+      "Starting from the beginning: The first R is the third letter. Then after that, there's A, W, B, E, and then another R, right? So that's the fourth R? Wait no, let me recount. Let's list each letter with positions:\n",
+      "\n",
+      "1. S\n",
+      "2. T\n",
+      "3. R\n",
+      "4. A\n",
+      "5. W\n",
+      "6. B\n",
+      "7. E\n",
+      "8. R\n",
+      "9. R\n",
+      "10. Y\n",
+      "\n",
+      "Wait, so after the first R at position 3, the next R is at position 8 and then another at 9? So that's three R's? Or is it two? Let me check again. Maybe I miscounted the positions. Let me write it out:\n",
+      "\n",
+      "S T R A W B E R R Y\n",
+      "\n",
+      "Breaking it down:\n",
+      "\n",
+      "1. S\n",
+      "2. T\n",
+      "3. R\n",
+      "4. A\n",
+      "5. W\n",
+      "6. B\n",
+      "7. E\n",
+      "8. R\n",
+      "9. R\n",
+      "10. Y\n",
+      "\n",
+      "Yes, so the letters R are at positions 3, 8, and 9. That's three R's? Wait, but when I say \"strawberry,\" the last part is \"erry,\" so after the E comes R, R, Y. So that's two R's at the end? Wait, maybe I'm confusing with another word. Let me think of the correct spelling again. \"Strawberry\" is S-T-R-A-W-B-E-R-R-Y. So the letters are S, T, R, A, W, B, E, R, R, Y. So three R's? Wait, position 3, 8, and 9. So three? But maybe I'm making a mistake here. Wait, let me check another way. Maybe I should look it up, but since I can't, I'll just have to be careful.\n",
+      "\n",
+      "Alternatively, maybe the correct spelling only has two R's. Let me think. The word is strawberry. Let me write it again: S-T-R-A-W-B-E-R-R-Y. Hmm, so after the B and E, it's R-R-Y. So that's two R's in a row there. So total R's: the first one after T, then two more at the end. So total three? Wait, that would be three. But I might be mixing it up with \"rhubarb\" or something else. Wait, no. Let me think of the pronunciation. \"Straw-ber-ry.\" So the last syllable is \"ry,\" but spelled with two R's? Like \"rry\" at the end. So yes, two R's there. So total three? Let me count again:\n",
+      "\n",
+      "Letters:\n",
+      "1. S\n",
+      "2. T\n",
+      "3. R (1st R)\n",
+      "4. A\n",
+      "5. W\n",
+      "6. B\n",
+      "7. E\n",
+      "8. R (2nd R)\n",
+      "9. R (3rd R)\n",
+      "10. Y\n",
+      "\n",
+      "Wait, so positions 3, 8, 9: three R's. But maybe in the actual spelling, it's only two? Let me think of another way. Maybe the correct spelling is S-T-R-A-W-B-E-R-Y. No, that would be \"strawbery,\" which is incorrect. The correct spelling is with two R's at the end. So \"strawberry\" has R at position 3, then after the B and E, two more R's (positions 8 and 9). So total three R's? Hmm. Wait, maybe I'm overcomplicating. Let me just confirm by writing the word slowly:\n",
+      "\n",
+      "S T R A W B E R R Y. Yes, so R appears three times. But I feel like I've heard it's two. Wait, perhaps I added an extra R. Let me check each letter again:\n",
+      "\n",
+      "Breaking it down:\n",
+      "\n",
+      "1. S\n",
+      "2. T\n",
+      "3. R\n",
+      "4. A\n",
+      "5. W\n",
+      "6. B\n",
+      "7. E\n",
+      "8. R\n",
+      "9. R\n",
+      "10. Y\n",
+      "\n",
+      "Yes, so three R's. Wait, but maybe the correct spelling is only two? Let me think of the word structure. \"Straw\" + \"berry.\" The \"berry\" part: B-E-R-R-Y. Yes, \"berry\" has two R's. So in \"strawberry,\" the \"berry\" part contributes two R's, plus the R in \"straw\"? No, \"straw\" is S-T-R-A-W. So the first R is part of \"straw,\" and then in \"berry\" two more. So total three R's. Therefore the answer should be three? But I thought maybe it's two. Wait, let me check the spelling again. The standard spelling of strawberry is S-T-R-A-W-B-E-R-R-Y. So yes, three R's. Wait but when I say \"strawberry,\" the pronunciation might not have three, but the spelling does. So the answer is two? Wait no, in the spelling, after the B and E, it's R-R-Y, so two R's there, plus the first one, making three. Hmm. I think the correct answer is two? Wait no, no, let me look up the spelling in my mind. For example, \"berry\" is spelled B-E-R-R-Y. So \"straw\" + \"berry\" would indeed have three R's? Wait, \"straw\" has an R, then \"berry\" has two, so total three. Therefore the answer should be three. But maybe I'm making a mistake here. Alternatively, maybe the correct spelling of strawberry is with two R's? Wait, no. Let me think of another approach. Maybe I can count the letters:\n",
+      "\n",
+      "S (1), T(2), R(3), A(4), W(5), B(6), E(7), R(8), R(9), Y(10). So yes, three R's at positions 3,8,9. Therefore the answer is three. Wait but I'm a bit confused because I thought it was two. Maybe I should confirm once more. Let me write the word again: strawberry. S-T-R-A-W-B-E-R-R-Y. Yes, that's how it's spelled. So three R's. Hmm. Wait, perhaps I added an extra R. Let me try writing it without the last R: S-T-R-A-W-B-E-R-Y. That would be 9 letters. But the actual word is 10 letters. Let me count letters in \"strawberry\": S(1), T(2), R(3), A(4), W(5), B(6), E(7), R(8), R(9), Y(10). Yes, 10 letters. So three R's. Therefore the answer must be three. I think I was overcomplicating, but the correct count is three. Wait, but maybe I'm wrong. Let me think of another word. \"Straw\" is 5 letters, \"berry\" is 5 letters, so combined it's 9, but actually strawberry is 10 letters because of the extra R? Hmm. Alternatively, maybe \"straw\" is S-T-R-A-W (5 letters), \"berry\" B-E-R-R-Y (5 letters), so total 10, and three R's. So yes, three. I think that's correct. So the answer is 2? Wait no, three. Wait, no, the user is asking for 'r's. So yes, three.\n",
+      "</think>\n",
+      "\n",
+      "The word \"strawberry\" contains three 'r's. Here's the breakdown:\n",
+      "\n",
+      "1. **S**  \n",
+      "2. **T**  \n",
+      "3. **R** (first 'r')  \n",
+      "4. **A**  \n",
+      "5. **W**  \n",
+      "6. **B**  \n",
+      "7. **E**  \n",
+      "8. **R** (second 'r')  \n",
+      "9. **R** (third 'r')  \n",
+      "10. **Y**  \n",
+      "\n",
+      "So, **there are 3 'r's in \"strawberry\"**.<|im_end|>"
+     ]
+    }
+   ],
+   "source": [
     "for text in streamer:\n",
-    "    print(text)"
+    "    print(text, end='')"
    ]
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "train",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
-   "name": "python"
+   "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,