From f61c9ab15fbfba0a4330820767931b461a15f368 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Fri, 2 May 2025 19:32:03 +0000
Subject: [PATCH] Add inference in notebook

---
 tools/notebook/gemma.ipynb | 91 ++++++++++++++++++++++++++++++++------
 1 file changed, 78 insertions(+), 13 deletions(-)

diff --git a/tools/notebook/gemma.ipynb b/tools/notebook/gemma.ipynb
index 8ce05be..24b6eaa 100644
--- a/tools/notebook/gemma.ipynb
+++ b/tools/notebook/gemma.ipynb
@@ -170,8 +170,9 @@
    "outputs": [],
    "source": [
     "# Set model ID\n",
-    "model_id = \"google/gemma-3-1b-it\"\n",
-    "#model_id = \"google/gemma-3-12b-it\""
+    "#model_id = \"google/gemma-3-1b-it\"\n",
+    "model_id = \"google/gemma-3-12b-it\"\n",
+    "output_dir = \"/root/models/notebook\""
    ]
   },
   {
@@ -286,7 +287,7 @@
     "    logging_steps=1,\n",
     "    save_strategy=\"steps\",\n",
     "    save_steps=1,\n",
-    "    output_dir=\"/root/models/notebook_lora\",\n",
+    "    output_dir=output_dir+\"_lora\",\n",
     "    optim=\"paged_adamw_8bit\",\n",
     "    seed=42,\n",
     "    group_by_length=True,\n",
@@ -321,6 +322,17 @@
     "trainer.train()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4e91a67e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Save tokenizer first\n",
+    "tokenizer.save_pretrained(output_dir+\"_merged\")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -329,7 +341,7 @@
    "outputs": [],
    "source": [
     "# Save the trained LoRA adapter\n",
-    "trainer.model.save_pretrained(\"/root/models/notebook_lora_adapter\")"
+    "trainer.model.save_pretrained(output_dir+\"_lora_adapter\")"
    ]
   },
   {
@@ -376,20 +388,73 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4e91a67e",
+   "id": "57b17d23",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Save tokenizer first\n",
-    "tokenizer.save_pretrained(\"/root/models/notebook_merged\")\n",
-    "\n",
     "# Save the merged model\n",
     "merged_model.save_pretrained(\n",
-    "    \"/root/models/notebook_merged\",\n",
+    "    output_dir+\"_merged\",\n",
     "    safe_serialization=True\n",
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0b1e39e1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import TextIteratorStreamer, pipeline\n",
+    "from threading import Thread\n",
+    "\n",
+    "messages = [\n",
+    "    {\"role\": \"user\", \"content\": \"Hi, write a poem about iced-tea.\"},\n",
+    "    {\"role\": \"assistant\", \"content\": \"\"},\n",
+    "]\n",
+    "prompt = tokenizer.apply_chat_template(\n",
+    "    messages, tokenize=False,\n",
+    "    add_generation_prompt=False,\n",
+    ")\n",
+    "pipeline = pipeline(\n",
+    "    \"text-generation\",\n",
+    "    model=merged_model,\n",
+    "    tokenizer=tokenizer,\n",
+    "    return_full_text=False,\n",
+    ")\n",
+    "streamer = TextIteratorStreamer(\n",
+    "    tokenizer,\n",
+    "    skip_prompt=True\n",
+    ")\n",
+    "generation_kwargs = {\n",
+    "    \"text_inputs\": prompt,\n",
+    "    \"do_sample\": True,\n",
+    "    \"max_new_tokens\": 1024,\n",
+    "    \"streamer\": streamer,\n",
+    "}\n",
+    "generation_thread = Thread(\n",
+    "    target=pipeline,\n",
+    "    kwargs=generation_kwargs\n",
+    ")\n",
+    "generation_thread.start()\n",
+    "\n",
+    "for text in streamer:\n",
+    "    print(text, end=\"\")\n",
+    "\n",
+    "generation_thread.join()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e5cc8064",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "exit(1)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -428,7 +493,7 @@
    "source": [
     "# Convert to GGUF (if merge was successful)\n",
     "!python ./llama.cpp/convert_hf_to_gguf.py \\\n",
-    "    --outfile /root/models/notebook.gguf \\\n",
+    "    --outfile /root/models/notebook_12B.gguf \\\n",
     "    --outtype q8_0 \\\n",
     "    /root/models/notebook_merged"
    ]
@@ -460,7 +525,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!unzip llama-b5226-bin-ubuntu-x64.zip"
+    "!unzip -o llama-b5226-bin-ubuntu-x64.zip"
    ]
   },
   {
@@ -476,9 +541,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "notebook",
    "language": "python",
-   "name": "python3"
+   "name": "notebook"
   },
   "language_info": {
    "codemirror_mode": {