examples/getting_started/olive-deepseek-finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tv6vx7wooDfk"
   },
   "source": [
    "# 🐋 Finetune and Optimize DeepSeek R1 with Olive\n",
    "\n",
    "DeepSeek introduced two first-generation reasoning models: DeepSeek-R1-Zero and DeepSeek-R1. DeepSeek-R1-Zero is trained using large-scale reinforcement learning (RL) without any supervised fine-tuning (SFT) as a preliminary step. This model has shown remarkable performance in reasoning tasks, naturally developing powerful and interesting reasoning behaviors through RL. However, it faces challenges such as endless repetition, poor readability, and language mixing.\n",
    "\n",
    "To address these issues and further enhance reasoning performance, DeepSeek developed DeepSeek-R1. This model incorporates cold-start data before applying RL, which helps mitigate the challenges faced by DeepSeek-R1-Zero. As a result, DeepSeek-R1 achieves performance comparable to OpenAI-o1 across various tasks, including math, code, and reasoning.\n",
    "\n",
    "To support the also released six dense models distilled from DeepSeek-R1. These dense models are based on Llama and Qwen architectures. In this notebook, you will:\n",
    "\n",
    "1. Fine-tune [DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B) for advanced medical reasoning.\n",
    "1. Optimize the fine-tuned model for the ONNX Runtime.\n",
    "1. Inference the fine-tuned model using the Generate API for the ONNX Runtime.\n",
    "\n",
    "## Motivation for Fine tuning DeepSeek R1\n",
    "\n",
    "Medical tasks frequently require intricate reasoning. In real-world scenarios, doctors often engage in careful deliberation when making diagnoses or decisions. To illustrate how a language model (LM) can assist in this process, consider a scenario where the LM receives a question about a patient. The model then employs chain-of-thought (CoT) reasoning to analyze the information and arrive at a diagnosis.\n",
    "\n",
    "Given a medical question, the distilled DeepSeek-R1-Distill-Qwen-1.5B base model can provide answers. However, they are overly verbose answers and not concise. For example, given the following question:\n",
    "\n",
    "*A 54-year-old construction worker with a long history of smoking presents with swelling in his upper extremity and face, along with dilated veins in this region. After conducting a CT scan and venogram of the neck, what is the most likely diagnosis for the cause of these symptoms?*\n",
    "\n",
    "The model responds with:\n",
    "\n",
    "*\\</think>*\n",
    "\n",
    "*Okay, so I'm trying to figure out what's going on with this 54-year-old construction worker who's been smoking and has these swelling and veins in his upper extremity and face. Let me break this down step by step. First, the age is 54, which is pretty old, but construction workers can be at higher risk for various health issues. The fact that he's been smoking for a long time suggests he might have a smoking-related condition, maybe something like lung cancer, but I'm not sure yet. He has swelling in his upper extremity and face, which is a common complaint in many patients, especially older ones. Swelling can be due to a variety of reasons, like infections, inflammation, or even something like a chronic condition. But given that he's been smoking, I'm thinking it's more likely something related to smoking. The CT scan and venogram results are important. CT scans can help visualize the areas, so the upper extremity and face would be checked. Venograms, or venous imaging, look at blood flow in the neck, which could indicate areas with high blood flow, possibly related to the lungs or other organs. Looking at the symptoms, the presence of blood vessels in that area suggests that there's an issue with blood flow, maybe an aneurysm or a blockage. An aneurysm is a bulge in a blood vessel, which can cause swelling and pain. But I'm not sure if that's the case here. Another possibility is atherosclerosis, where plaque builds up in the blood vessels, leading to blockage and swelling. This is common in older patients, especially those who smoke, because smoking increases the risk of lung cancer, which can lead to atherosclerosis. I should also consider other conditions, like a chronic disease like diabetes, but the question doesn't mention that. The main focus seems to be on smoking-related issues. Putting it all together, the CT scan and venogram likely showed areas with high blood flow, which points to a blockage in the blood vessels in the upper extremity and face. This could be an aneurysm or atherosclerosis. Given that he's been smoking, atherosclerosis is a more likely culprit, but aneurysm could also be a possibility. I should also think about the patient's smoking history. If he's been smoking for a long time, there's a higher risk of lung cancer, which can lead to atherosclerosis. Smoking also increases the risk of other conditions like chronic obstructive pulmonary disease (COPD), which can cause blockages.So, the most likely diagnosis would be a blockage in the blood vessels, either an aneurysm or atherosclerosis. But without more specific imaging, it's hard to tell exactly. However, given the smoking history, atherosclerosis is a strong candidate. I should also consider other possibilities, like a tumor, but that's less likely given the symptoms. The main suspects are atherosclerosis or aneurysm due to smoking-related conditions. In conclusion, the most likely diagnosis is a blockage in the blood vessels in the upper extremity and face, likely due to smoking-related atherosclerosis or a nearby aneurysm.*\n",
    "\n",
    "*\\</think>*\n",
    "\n",
    "*The most likely diagnosis for the described symptoms is a blockage in the blood vessels in the upper extremity and face, likely due to smoking-related atherosclerosis or a nearby aneurysm. This conclusion is based on the patient's smoking history, which increases the risk of lung cancer and atherosclerosis, both of which can lead to blood vessel blockages.*\n",
    "\n",
    "The reasoning process is outlined within the \\<think>\\</think> tags. We'd prefer to have more concise reasoning. To this end, we'll finetune the DeepSeek-R1-Distill-Qwen-1.5B model using the [FreedomIntelligence/medical-o1-reasoning-SFT](https://huggingface.co/datasets/FreedomIntelligence/medical-o1-reasoning-SFT) dataset, which was constructed using GPT-4o and searches for solutions to [verifiable medical problems](https://huggingface.co/datasets/FreedomIntelligence/medical-o1-verifiable-problem). The data is validated through a medical verifier. The dataset comprises of 3 fields:\n",
    "\n",
    "1. **Question**: A medical questions, explaining patient symptoms and background.\n",
    "2. **Complex_COT**: The Chain of Thought (CoT) the LM should go through given the question.\n",
    "3. **Response**: The formal response in the form of a diagnosis of the condition with a brief rationale.\n",
    "\n",
    "The following is a single sample taken from the dataset:\n",
    "\n",
    "| Question | Complex_COT | Response |\n",
    "|----------| ----------- | ---------|\n",
    "|A 17-year-old girl presents with sudden onset left-sided facial droop and is unable to fully close her left eye. Given her symptoms of unilateral facial paralysis that include an inability to move the left side of her mouth and close her left eye, but with intact extraocular movements and no other neurologic deficits, what treatment is most likely to address and improve her condition?| Alright, so we've got a 17-year-old girl here who suddenly can't move the left side of her face and can't fully close her left eye. That's pretty concerning. When I hear about facial droop and eye issues like that, I'm thinking about facial nerve stuff, like Bell's palsy. Okay, let's break this down. She's young and it's sudden onset, which fits the typical presentation of Bell's palsy. Plus, she's only got the one-sided facial paralysis, and no other neurological red flags—so things like her eye movements are normal. That's important because it helps to make sure this isn't something like a stroke or another cranial nerve problem. Hmm, definitely sounds like we're on the right track with Bell's palsy. Now, what do we do for Bell's palsy again? Oh right, corticosteroids! Those are the go-to treatment. Prednisone is what normally comes up because it helps reduce inflammation around the facial nerve. And you want to start that early, like ideally within the first three days or so of the symptoms showing up. So timing is key here. Also, we've got to think about her eye. If she can't close it properly, she might run into issues with dryness or injury. So we need to make sure she protects that eye. Keeping it moist with artificial tears or some kind of eye lubricant would be necessary. Putting it all together, for her condition, starting her on corticosteroids seems like the right move. Plus, taking care of her eye is crucial to avoid any complications there. This combo should set her on the path to recovery. Everything we've looked at lines up, and this plan makes sense based on her symptoms and the expected treatment for Bell's palsy.|Based on the symptoms you're describing, the 17-year-old girl is most likely experiencing Bell's palsy, which is characterized by the sudden onset of unilateral facial paralysis. The primary treatment for Bell's palsy is corticosteroids, with prednisone being the commonly used medication. Administering steroids can help reduce inflammation and swelling of the facial nerve, promoting recovery of nerve function. It is most effective when started early, ideally within the first three days of symptom onset. Additionally, it's crucial to manage her eye care, as she cannot fully close her left eye. Using artificial tears or eye lubricants will help keep her eye moist and protected from potential injury or dryness due to incomplete closure. In summary, the treatment that is most likely to address and improve her condition involves starting corticosteroids, such as prednisone, and ensuring protective care for her affected eye.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🐍 Install Python dependencies\n",
    "\n",
    "The following cells create a pip requirements file and then install the libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile requirements.txt\n",
    "\n",
    "olive-ai\n",
    "onnxruntime-genai==0.6.0\n",
    "transformers==4.44.2\n",
    "optimum==1.23.1\n",
    "bitsandbytes==0.45.1\n",
    "peft==0.13.2\n",
    "accelerate>=0.30.0\n",
    "scipy==1.14.1\n",
    "torchvision==0.18.1\n",
    "tabulate==0.9.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZtY3VYxCoDfm"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ❔💭💬 Pre-process the data\n",
    "\n",
    "We need to massage the data so that a prompt is constructed for DeepSeek containing the instruction, medical question, human Chain of Thought (CoT) and desired response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "prompt_template = \"\"\"Below is an instruction that describes a task, paired with an input that provides further context. \n",
    "Write a response that appropriately completes the request. \n",
    "Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.\n",
    "\n",
    "### Instruction:\n",
    "You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning. \n",
    "Please answer the following medical question. \n",
    "\n",
    "### Question:\n",
    "{}\n",
    "\n",
    "### Response:\n",
    "<think>\n",
    "{}\n",
    "</think>\n",
    "{}\"\"\"\n",
    "\n",
    "def formatting_prompts_func(examples):\n",
    "    inputs = examples[\"Question\"]\n",
    "    cots = examples[\"Complex_CoT\"]\n",
    "    outputs = examples[\"Response\"]\n",
    "    texts = []\n",
    "    for input, cot, output in zip(inputs, cots, outputs):\n",
    "        text = prompt_template.format(input, cot, output) + \"<｜end▁of▁sentence｜>\"\n",
    "        texts.append(text)\n",
    "    return {\n",
    "        \"text\": texts,\n",
    "    }\n",
    "\n",
    "# Create the English dataset\n",
    "dataset = load_dataset(\"FreedomIntelligence/medical-o1-reasoning-SFT\",\"en\", split = \"train\",trust_remote_code=True)\n",
    "dataset = dataset.map(formatting_prompts_func, batched = True,remove_columns=[\"Question\", \"Complex_CoT\", \"Response\"])\n",
    "dataset.to_json(\"en_dataset.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🏃 Train the model\n",
    "\n",
    "In this tutorial we will only fine-tune for 100 steps, which minimizes the time to complete the tutorial. To improve the accuracy of the model, more training steps would be required. However, even with just 100 steps you'll notice later in the tutorial that the fine-tuned model is much more concise than using the base model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8t36pRF2oDfq"
   },
   "outputs": [],
   "source": [
    "!olive finetune \\\n",
    "    --method lora \\\n",
    "    --model_name_or_path \"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B\" \\\n",
    "    --trust_remote_code \\\n",
    "    --data_name json \\\n",
    "    --data_files ./en_dataset.jsonl \\\n",
    "    --train_split \"train[:20000]\" \\\n",
    "    --eval_split \"train[20000:25400]\" \\\n",
    "    --text_field \"text\" \\\n",
    "    --max_steps 100 \\\n",
    "    --logging_steps 10 \\\n",
    "    --target_modules \"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\"gate_proj\",\"up_proj\",\"down_proj\" \\\n",
    "    --output_path models/deepseek/en_ft \\\n",
    "    --log_level 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7woNXLDF0bhh"
   },
   "source": [
    "## 🪄 Automatic model optimization with Olive\n",
    "\n",
    "Next, you'll execute Olive's automatic optimizer using the `auto-opt` CLI command, which will:\n",
    "\n",
    "1. Capture the fine-tuned model into an ONNX graph and convert the weights into the ONNX format.\n",
    "1. Optimize the ONNX graph (e.g. fuse nodes, reshape, etc).\n",
    "1. Extract the fine-tuned LoRA weights and place them into a separate file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!olive auto-opt \\\n",
    "    --model_name_or_path \"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B\" \\\n",
    "    --adapter_path models/deepseek/en_ft/adapter \\\n",
    "    --device cpu \\\n",
    "    --provider CPUExecutionProvider \\\n",
    "    --use_ort_genai \\\n",
    "    --precision int4 \\\n",
    "    --output_path models/deepseek/en-onnx-ao \\\n",
    "    --log_level 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8Uwm432loDfr"
   },
   "source": [
    "## 🧠 Inference\n",
    "\n",
    "The code below creates a test app that consumes the finetuned model in a simple console chat interface. Whilst the inference code uses the Python API for the ONNX Runtime, other language bindings are available in [Java, C#, C++](https://github.com/microsoft/onnxruntime-genai/tree/main/examples).\n",
    "\n",
    "By default, the code will use the fine-tuned adapter to deliver more concise responses. To see the output of the original model (no fintuning), set `USE_ADAPTER = False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "puMdoAxjoDfr"
   },
   "outputs": [],
   "source": [
    "import onnxruntime_genai as og\n",
    "\n",
    "USE_ADAPTER = True\n",
    "\n",
    "model_path = \"models/deepseek/en-onnx-ao/model\"\n",
    "\n",
    "model = og.Model(f'{model_path}')\n",
    "adapters = og.Adapters(model)\n",
    "adapters.load(f'{model_path}/adapter_weights.onnx_adapter', \"en_medical_reasoning\")\n",
    "tokenizer = og.Tokenizer(model)\n",
    "tokenizer_stream = tokenizer.create_stream()\n",
    "\n",
    "prompt_template = \"\"\"\n",
    "Below is an instruction that describes a task, paired with an input that provides further context. \n",
    "Write a response that appropriately completes the request. \n",
    "Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.\n",
    "\n",
    "### Instruction:\n",
    "You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning. \n",
    "Please answer the following medical question. \n",
    "\n",
    "### Question:\n",
    "{}\n",
    "\n",
    "### Response:\n",
    "<think>\n",
    "\"\"\"\n",
    "\n",
    "question = \"\"\"\n",
    "        A 54-year-old construction worker with a long history of smoking presents with swelling in his upper extremity and face, along with \n",
    "        dilated veins in this region. After conducting a CT scan and venogram of the neck, what is the most likely diagnosis for the cause of these symptoms?\n",
    "\"\"\"\n",
    "prompt = prompt_template.format(question, \"\")\n",
    "input_tokens = tokenizer.encode(prompt)\n",
    "\n",
    "# first run without the adapter\n",
    "params = og.GeneratorParams(model)\n",
    "params.set_search_options(past_present_share_buffer=False, temperature=0.6, max_length=1200)\n",
    "generator = og.Generator(model, params)\n",
    "# set the adapter to active for this response\n",
    "\n",
    "if USE_ADAPTER:\n",
    "    generator.set_active_adapter(adapters, \"en_medical_reasoning\")\n",
    "\n",
    "generator.append_tokens(input_tokens)\n",
    "\n",
    "print(\"Output: \", end='', flush=True)\n",
    "\n",
    "while not generator.is_done():\n",
    "    generator.generate_next_token()\n",
    "    new_token = generator.get_next_tokens()[0]\n",
    "    print(tokenizer_stream.decode(new_token), end='', flush=True)\n",
    "print()\n",
    "print()\n",
    "del generator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 💬 Output Before/After\n",
    "\n",
    "The output of Deepseek R1 before/after finetuning is shown in the following table - notice that the output from the fine-tuned model is more concise than the base model. Whilst the conciseness has improved, more training would be required to improve the quality of the outputs.\n",
    "\n",
    "|**Before Finetuning<br>(550 tokens generated)**|**After Finetuning<br>(455 tokens generated)**|\n",
    "|:---------------------|:--------------------|\n",
    "|*Okay, so I'm trying to figure out what's going on with this 54-year-old construction worker who's been smoking and has these swelling and veins in his upper extremity and face. Let me break this down step by step.<br><br>First, the age is 54, which is pretty old, but construction workers can be at higher risk for various health issues. The fact that he's been smoking for a long time suggests he might have a smoking-related condition, maybe something like lung cancer, but I'm not sure yet.<br><br>He has swelling in his upper extremity and face, which is a common complaint in many patients, especially older ones. Swelling can be due to a variety of reasons, like infections, inflammation, or even something like a chronic condition. But given that he's been smoking, I'm thinking it's more likely something related to smoking.<br><br>The CT scan and venogram results are important. CT scans can help visualize the areas, so the upper extremity and face would be checked. Venograms, or venous imaging, look at blood flow in the neck, which could indicate areas with high blood flow, possibly related to the lungs or other organs.<br><br>Looking at the symptoms, the presence of blood vessels in that area suggests that there's an issue with blood flow, maybe an aneurysm or a blockage. An aneurysm is a bulge in a blood vessel, which can cause swelling and pain. But I'm not sure if that's the case here.<br><br>Another possibility is atherosclerosis, where plaque builds up in the blood vessels, leading to blockage and swelling. This is common in older patients, especially those who smoke, because smoking increases the risk of lung cancer, which can lead to atherosclerosis.<br><br>I should also consider other conditions, like a chronic disease like diabetes, but the question doesn't mention that. The main focus seems to be on smoking-related issues.<br><br>Putting it all together, the CT scan and venogram likely showed areas with high blood flow, which points to a blockage in the blood vessels in the upper extremity and face. This could be an aneurysm or atherosclerosis. Given that he's been smoking, atherosclerosis is a more likely culprit, but aneurysm could also be a possibility.<br><br>I should also think about the patient's smoking history. If he's been smoking for a long time, there's a higher risk of lung cancer, which can lead to atherosclerosis. Smoking also increases the risk of other conditions like chronic obstructive pulmonary disease (COPD), which can cause blockages.<br><br>So, the most likely diagnosis would be a blockage in the blood vessels, either an aneurysm or atherosclerosis. But without more specific imaging, it's hard to tell exactly. However, given the smoking history, atherosclerosis is a strong candidate.<br><br>I should also consider other possibilities, like a tumor, but that's less likely given the symptoms. The main suspects are atherosclerosis or aneurysm due to smoking-related conditions.<br><br>In conclusion, the most likely diagnosis is a blockage in the blood vessels in the upper extremity and face, likely due to smoking-related atherosclerosis or a nearby aneurysm.<br>\\</think><br><br>The most likely diagnosis for the described symptoms is a blockage in the blood vessels in the upper extremity and face, likely due to smoking-related atherosclerosis or a nearby aneurysm. This conclusion is based on the patient's smoking history, which increases the risk of lung cancer and atherosclerosis, both of which can lead to blood vessel blockages.*|*Okay, so we've got a 54-year-old construction worker who's been smoking for a while. That's a big hint that he might be dealing with something related to his lung health. Now, he's showing up with swelling in his upper extremity and face, and there are these dilated veins in that area too. Hmm, that's interesting.<br><br>First off, I'm thinking about what could cause these symptoms. Swelling in the neck, especially in the upper extremity, could be due to a few things. Maybe it's something like a tumor, like a lung cancer, which can cause swelling. Or perhaps it's something more subtle, like a tumor in the brain, which might cause a lot of swelling too.<br><br>Then there's the venous dilation. That's a clue too. If there's a tumor in the neck, it could cause the blood vessels to dilate, which would explain the swelling. But wait, there's also the possibility of a tumor in the brain, which could cause a lot of swelling too, but it's more about the brain's blood supply rather than the neck.<br><br>Now, let's think about what's more likely. If we're talking about a tumor in the neck, it's more likely to cause the neck's swelling, and it's also more likely to cause the venous dilation. That's because tumors in the neck are more likely to be in the brain, but they can also be in the neck itself.<br><br>Oh, and there's also the fact that he's been smoking for a long time. Smoking can lead to lung cancer, which is a common cause of lung tumors. So, a lung cancer tumor in the neck could explain both the neck swelling and the venous dilation.<br><br>But wait, could there be something else going on? Maybe a tumor in the brain could cause similar symptoms, but it's less likely to cause the neck's swelling as much as a tumor in the neck itself.<br><br>So, putting it all together, I'm leaning towards a lung cancer tumor in the neck. It's the most likely cause of both the neck swelling and the venous dilation. Yeah, that makes sense. It's a pretty common scenario in lung cancer cases.<br><br>Yeah, I'm pretty confident now that it's a lung cancer tumor in the neck. That's the most likely diagnosis here.<br>\\</think><br><br>The most likely diagnosis for the symptoms described in the question is a lung cancer tumor in the neck. This condition, known as a \"neoplasmic pheochromocytoma,\" is a common cause of both neck swelling and venous dilation. The symptoms of swelling in the upper extremity and face, along with the dilation of venous vessels, are consistent with a tumor in the neck, which can be associated with a tumor in the brain due to the possibility of a brain tumor causing similar systemic effects.*<br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br>|"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "genai-cpu",
   "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}