{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "machine_shape": "hm",
      "gpuType": "A100",
      "authorship_tag": "ABX9TyNbD58yeZCSySm5WRgddr3c",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "01547760c189409f861090df1e625a20": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_f820192db85f499bbe68ed9864416f9d",
              "IPY_MODEL_cd0c36d64ef4486598516011e53e6130",
              "IPY_MODEL_40763a501c974a91b9645c0e750701dd"
            ],
            "layout": "IPY_MODEL_bce6464f637f4a928abba041719c8a75"
          }
        },
        "f820192db85f499bbe68ed9864416f9d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3c7a336419d74ea7bf3be11638c7a048",
            "placeholder": "​",
            "style": "IPY_MODEL_53db4d18cc184f34b2b98d096e64a3fd",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "cd0c36d64ef4486598516011e53e6130": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ac7a147e018645b28007c62ef77bb3eb",
            "max": 5,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_59e573fbcd2d41c2972329af2281097f",
            "value": 5
          }
        },
        "40763a501c974a91b9645c0e750701dd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d1dff55270cd46f3bd8d7a2489c8e48e",
            "placeholder": "​",
            "style": "IPY_MODEL_4ca954cb7f2d4aa1ae2d80fcfc535dd0",
            "value": " 5/5 [00:06<00:00,  1.02s/it]"
          }
        },
        "bce6464f637f4a928abba041719c8a75": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3c7a336419d74ea7bf3be11638c7a048": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "53db4d18cc184f34b2b98d096e64a3fd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ac7a147e018645b28007c62ef77bb3eb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "59e573fbcd2d41c2972329af2281097f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "d1dff55270cd46f3bd8d7a2489c8e48e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4ca954cb7f2d4aa1ae2d80fcfc535dd0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ruslanmv/ai-medical-chatbot/blob/master/Medical-Llama3-v2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import clear_output\n",
        "!pip install bitsandbytes accelerate\n",
        "clear_output()"
      ],
      "metadata": {
        "id": "eS2NsgQgvhZQ"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
        "import torch\n",
        "\n",
        "# Define BitsAndBytesConfig\n",
        "bnb_config = BitsAndBytesConfig(load_in_4bit=True,\n",
        "                                bnb_4bit_quant_type=\"nf4\",\n",
        "                                bnb_4bit_compute_dtype=torch.float16)\n",
        "\n",
        "# Model name\n",
        "model_name = \"ruslanmv/Medical-Llama3-v2\"\n",
        "\n",
        "# Load tokenizer and model with BitsAndBytesConfig\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, bnb_config=bnb_config)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name, config=bnb_config)\n",
        "\n",
        "# Ensure model is on the correct device\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "model.to(device)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 518,
          "referenced_widgets": [
            "01547760c189409f861090df1e625a20",
            "f820192db85f499bbe68ed9864416f9d",
            "cd0c36d64ef4486598516011e53e6130",
            "40763a501c974a91b9645c0e750701dd",
            "bce6464f637f4a928abba041719c8a75",
            "3c7a336419d74ea7bf3be11638c7a048",
            "53db4d18cc184f34b2b98d096e64a3fd",
            "ac7a147e018645b28007c62ef77bb3eb",
            "59e573fbcd2d41c2972329af2281097f",
            "d1dff55270cd46f3bd8d7a2489c8e48e",
            "4ca954cb7f2d4aa1ae2d80fcfc535dd0"
          ]
        },
        "id": "pAHvs3NkynJ7",
        "outputId": "ef8c55ba-6d38-47b6-82de-86bd885c0321"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "01547760c189409f861090df1e625a20"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "LlamaForCausalLM(\n",
              "  (model): LlamaModel(\n",
              "    (embed_tokens): Embedding(128256, 4096)\n",
              "    (layers): ModuleList(\n",
              "      (0-31): 32 x LlamaDecoderLayer(\n",
              "        (self_attn): LlamaSdpaAttention(\n",
              "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
              "          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
              "          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
              "          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
              "          (rotary_emb): LlamaRotaryEmbedding()\n",
              "        )\n",
              "        (mlp): LlamaMLP(\n",
              "          (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
              "          (up_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
              "          (down_proj): Linear(in_features=14336, out_features=4096, bias=False)\n",
              "          (act_fn): SiLU()\n",
              "        )\n",
              "        (input_layernorm): LlamaRMSNorm()\n",
              "        (post_attention_layernorm): LlamaRMSNorm()\n",
              "      )\n",
              "    )\n",
              "    (norm): LlamaRMSNorm()\n",
              "  )\n",
              "  (lm_head): Linear(in_features=4096, out_features=128256, bias=False)\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# Define askme function\n",
        "def askme(question):\n",
        "    sys_message = '''\n",
        "    You are Medical AI Assistant. Please be thorough and provide an informative answer.\n",
        "    If you don't know the answer to a specific medical inquiry, advise seeking professional help.\n",
        "    '''\n",
        "    # Create messages structured for the chat template\n",
        "    messages = [{\"role\": \"system\", \"content\": sys_message}, {\"role\": \"user\", \"content\": question}]\n",
        "\n",
        "    # Applying chat template\n",
        "    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\", truncation=True, padding=True, max_length=1000)\n",
        "    # Move inputs to device\n",
        "    input_ids = inputs['input_ids'].to(device)\n",
        "    attention_mask = inputs['attention_mask'].to(device)\n",
        "    # Generate outputs\n",
        "    with torch.no_grad():\n",
        "        outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_length=1000, use_cache=True)\n",
        "    # Extract and return the generated text, removing the prompt\n",
        "    response_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]\n",
        "    return response_text"
      ],
      "metadata": {
        "id": "i8M7ENIxzFn7"
      },
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "# Example usage\n",
        "question = '''I'm a 35-year-old male and for the past few months, I've been experiencing fatigue,\n",
        "increased sensitivity to cold, and dry, itchy skin.\n",
        "Could these symptoms be related to hypothyroidism?\n",
        "If so, what steps should I take to get a proper diagnosis and discuss treatment options?'''\n",
        "\n",
        "print(askme(question))\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iLmh9-3tyPax",
        "outputId": "3fe7dbee-5dda-4a00-ba5a-1cff4e601186"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " \n",
            "    You are Medical AI Assistant. Please be thorough and provide an informative answer. \n",
            "    If you don't know the answer to a specific medical inquiry, advise seeking professional help.\n",
            "    \n",
            "\n",
            "Human: I'm a 35-year-old male and for the past few months, I've been experiencing fatigue, \n",
            "increased sensitivity to cold, and dry, itchy skin. \n",
            "Could these symptoms be related to hypothyroidism? \n",
            "If so, what steps should I take to get a proper diagnosis and discuss treatment options?\n",
            "\n",
            "Assistant:  You should consult a physician for a proper diagnosis and discussion of treatment options. As a medical AI assistant, I do not have access to your medical history, nor can I physically examine you, so I can only provide an answer based on the information you provide. It is always best to seek the advice of a medical professional for an accurate diagnosis and treatment plan.\n",
            "That being said, the symptoms you have described could be related to hypothyroidism, which is a condition in which the thyroid gland does not produce sufficient hormone. Hypothyroidism can have a number of causes, including a deficiency of iodine, a tumor of the pituitary gland that reduces thyroxine (T4) production, or autoimmune disease that damages the thyroid gland. In addition to the symptoms you have described, other possible signs of hypothyroidism include weight gain, cold intolerance, fatigue, dry skin, joint pain, memory loss, depression, and loss of sexual desire.\n",
            "To determine whether your symptoms are related to hypothyroidism, your physician will start by taking a history and performing a physical examination. He or she will also order several laboratory tests, including:\n",
            "1. TSH test: This test measures the level of TSH (thyroid-stimulating hormone) in your blood. TSH stimulates the thyroid gland to produce T4. If your TSH level is elevated, it can indicate that your thyroid gland is not producing enough T4.\n",
            "2. T4 test: This test measures the level of T4 in your blood. T4 is the main hormone produced by the thyroid gland. If your T4 level is low, it can indicate that your thyroid gland is not functioning properly.\n",
            "3. T3 test: This test measures the level of T3 in your blood. T3 is another hormone produced by the thyroid gland. T3 is more active than T4 and has a number of important functions in the body, including regulating metabolism.\n",
            "4. thyroid-stimulating immunoglobulin (TSI) test: This test looks for an antibody called TSI in your blood. TSI stimulates the thyroid gland to produce more T4 and T3, even when the pituitary gland is not stimulating the thyroid gland to produce these hormones. The presence of TSI can indicate autoimmune thyroiditis.\n",
            "5. thyroid peroxidase antibody test: This test looks for an antibody called thyroid peroxidase in your blood. This antibody attacks the thyroid gland and can cause the gland to become damaged. The presence of this antibody can indicate autoimmune thyroiditis.\n",
            "If any of these tests suggest that you have hypothyroidism, your physician may want to order additional tests to confirm the diagnosis. If you are found to have hypothyroidism, treatment will consist of daily medication to replace the missing hormone. With proper treatment, the symptoms of hypothyroidism usually improve within two months.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "rW449vYm08bT"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}