MedQA_Llama3_1_8B_RAG.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "A100"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7cjF2e2EYeqG",
        "outputId": "f3a15cd2-efa7-45dd-b791-b85cbce5c853"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n",
            "==((====))==  Unsloth 2025.1.5: Fast Llama patching. Transformers: 4.47.1.\n",
            "   \\\\   /|    GPU: NVIDIA A100-SXM4-40GB. Max memory: 39.564 GB. Platform: Linux.\n",
            "O^O/ \\_/ \\    Torch: 2.5.1+cu121. CUDA: 8.0. CUDA Toolkit: 12.1. Triton: 3.1.0\n",
            "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.29.post1. FA2 = False]\n",
            " \"-____-\"     Free Apache license: http://github.com/unslothai/unsloth\n",
            "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n",
            "Question: What is the role of insulin in managing diabetes?\n",
            "Answer: Insulin is the only treatment for diabetes that can normalize both blood glucose and blood pressure. It is also the only treatment that can prevent microvascular complications. In this context, insulin is a first-line treatment for type 2 diabetes. It is also indicated as an adjunct to oral antidiabetic drugs in type 1 diabetes and type 2 diabetes. Insulin therapy is indicated in all cases of diabetes with ketoacidosis, hypoglycemia, or hyperosmolar coma. It is also indicated in type 1 diabetes, type 2 diabetes, and gestational diabetes when blood glucose is poorly controlled despite optimal oral antidiabetic therapy. Insulin therapy should be considered in patients with type 2 diabetes and a hemoglobin A1c (HbA1c) level >or = 8.0% despite optimal oral antidiabetic therapy. Insulin is also indicated in type 2 diabetes when the patient's HbA1c level is >or = 9.0% despite optimal oral antidiabetic therapy and a low risk of hypoglycemia. Insulin therapy should be initiated in type 2 diabetes when the patient's HbA1c level is >or = 10.0% despite optimal oral\n"
          ]
        }
      ],
      "source": [
        "!pip install unsloth datasets langchain faiss-gpu transformers sentence-transformers langchain-community -q\n",
        "\n",
        "import torch\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "drive_path = \"/content/drive/My Drive/MedQA-Llama3.1-8B_LoRA_Model/lora_model\"\n",
        "\n",
        "from unsloth import FastLanguageModel\n",
        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
        "    model_name=drive_path,\n",
        "    max_seq_length=2048,\n",
        "    dtype=None,\n",
        "    load_in_4bit=True,\n",
        "    device_map=\"auto\"\n",
        ")\n",
        "FastLanguageModel.for_inference(model)\n",
        "\n",
        "from datasets import load_dataset\n",
        "dataset = load_dataset(\"bigbio/pubmed_qa\", name=\"pubmed_qa_labeled_fold0_source\", split=\"train\", trust_remote_code=True)\n",
        "docs = []\n",
        "for i, ex in enumerate(dataset):\n",
        "    t = ex[\"LONG_ANSWER\"] or \"\"\n",
        "    docs.append({\"id\": f\"pubmed_qa_train_{i}\", \"text\": t})\n",
        "\n",
        "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
        "ts = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)\n",
        "chunks = []\n",
        "for d in docs:\n",
        "    c = ts.split_text(d[\"text\"])\n",
        "    for idx, x in enumerate(c):\n",
        "        chunks.append({\"id\": f\"{d['id']}_chunk_{idx}\", \"text\": x})\n",
        "\n",
        "from langchain.embeddings import HuggingFaceEmbeddings\n",
        "from langchain.vectorstores import FAISS\n",
        "e = HuggingFaceEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\")\n",
        "texts = [x[\"text\"] for x in chunks]\n",
        "meta = [{\"source\": x[\"id\"]} for x in chunks]\n",
        "db = FAISS.from_texts(texts, embedding=e, metadatas=meta)\n",
        "\n",
        "from pydantic import PrivateAttr\n",
        "from langchain.llms.base import LLM\n",
        "from typing import Any, Optional, List\n",
        "\n",
        "class LoRAMedicalLLM(LLM):\n",
        "    _model: Any = PrivateAttr()\n",
        "    _tokenizer: Any = PrivateAttr()\n",
        "    max_new_tokens: int = 512\n",
        "    device: str = \"cuda\"\n",
        "    def __init__(self, model, tokenizer, max_new_tokens=256, device=\"cuda\", **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self._model = model\n",
        "        self._tokenizer = tokenizer\n",
        "        self.max_new_tokens = max_new_tokens\n",
        "        self.device = device\n",
        "    @property\n",
        "    def _llm_type(self) -> str:\n",
        "        return \"lora-medical-llm\"\n",
        "    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:\n",
        "        inp = self._tokenizer([prompt], return_tensors=\"pt\", padding=True, truncation=True, max_length=2048).to(self.device)\n",
        "        with torch.no_grad():\n",
        "            out_toks = self._model.generate(**inp, max_new_tokens=self.max_new_tokens, use_cache=True)\n",
        "        return self._tokenizer.decode(out_toks[0], skip_special_tokens=True)\n",
        "\n",
        "llm = LoRAMedicalLLM(model=model, tokenizer=tokenizer)\n",
        "\n",
        "from langchain.prompts import PromptTemplate\n",
        "from langchain.chains import RetrievalQA\n",
        "pt = \"\"\"You are a medical QA system.\n",
        "Use the following context to answer the question concisely, dont repeat the context or any lines in the final answer:\n",
        "{context}\n",
        "\n",
        "Question: {question}\n",
        "\n",
        "Answer:\n",
        "\"\"\"\n",
        "pr = PromptTemplate(template=pt, input_variables=[\"context\",\"question\"])\n",
        "retriever = db.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 3})\n",
        "rag = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\", retriever=retriever, chain_type_kwargs={\"prompt\":pr})\n",
        "\n",
        "q = \"What is the role of insulin in managing diabetes?\"\n",
        "r = rag({\"query\": q})\n",
        "raw_answer = r[\"result\"]\n",
        "ans = raw_answer.split(\"Answer:\")[-1].strip() if \"Answer:\" in raw_answer else raw_answer.strip()\n",
        "\n",
        "print(\"Question:\", q)\n",
        "print(\"Answer:\", ans)\n"
      ]
    }
  ]
}