Optimize_SBERT/BAAI_bge_base_zh.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cac7597b-2177-47f8-884f-b4b4fd6598ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "from optimum.onnxruntime import (\n",
    "    AutoOptimizationConfig,\n",
    "    ORTModelForFeatureExtraction,\n",
    "    ORTOptimizer,\n",
    "    ORTQuantizer\n",
    ")\n",
    "from optimum.onnxruntime.configuration import OptimizationConfig, AutoQuantizationConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d22b37d2-a75a-4ba2-a34e-14ecdfa7d31d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the tokenizer and export the model to the ONNX format\n",
    "model_id = \"BAAI/bge-base-zh\"\n",
    "save_dir = \"BAAI/bge-base-zh_optimized\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c4941e-08c2-4a28-89e9-a5f838d403e4",
   "metadata": {},
   "source": [
    "## Optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "07d142ee-2459-40d4-9301-770deb62977c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4b5a91f00d674262b514bfc0c47325a8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/366 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8518e4eed2104bde8750d9e1db925b49",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.txt:   0%|          | 0.00/110k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2567f10343ce48c3bd309d048710d1ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/439k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a112ae75f3a94e4a98883a38665817ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/125 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "92efe4ada3214219ae2359c9e7fc2e83",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/940 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Framework not specified. Using pt to export the model.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c738c7802be498bb4c9c4f8e573d42d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/409M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using the export variant default. Available variants are:\n",
      "    - default: The default ONNX variant.\n",
      "Using framework PyTorch: 2.2.0\n",
      "Overriding 1 configuration item(s)\n",
      "\t- use_cache -> False\n",
      "/Users/satishsilveri/Documents/ML/optimum/optimum/onnxruntime/configuration.py:770: FutureWarning: disable_embed_layer_norm will be deprecated soon, use disable_embed_layer_norm_fusion instead, disable_embed_layer_norm_fusion is set to True.\n",
      "  warnings.warn(\n",
      "Optimizing model...\n",
      "Configuration saved in BAAI/bge-base-zh_optimized/ort_config.json\n",
      "Optimized model saved at: BAAI/bge-base-zh_optimized (external data format: False; saved all tensor to one file: True)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PosixPath('BAAI/bge-base-zh_optimized')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "model = ORTModelForFeatureExtraction.from_pretrained(model_id, export=True)\n",
    "\n",
    "# Load the optimization configuration detailing the optimization we wish to apply\n",
    "optimization_config = AutoOptimizationConfig.O3()\n",
    "optimizer = ORTOptimizer.from_pretrained(model)\n",
    "\n",
    "optimizer.optimize(save_dir=save_dir, optimization_config=optimization_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a211542-1d4e-482e-a395-a0c6b97dcc63",
   "metadata": {},
   "source": [
    "## Quantize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d2fb3ecd-2993-4be0-b687-6d7007399d1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using a model of type bert to instantiate a model of type . This is not supported for all configurations of models and can yield errors.\n",
      "Creating dynamic quantizer: QOperator (mode: IntegerOps, schema: u8/s8, channel-wise: False)\n",
      "Quantizing model...\n",
      "Saving quantized model at: BAAI/bge-base-zh_optimized_quantized (external data format: False)\n",
      "Configuration saved in BAAI/bge-base-zh_optimized_quantized/ort_config.json\n"
     ]
    }
   ],
   "source": [
    "onnx_model = ORTModelForFeatureExtraction.from_pretrained(save_dir)\n",
    "\n",
    "quantizer = ORTQuantizer.from_pretrained(onnx_model)\n",
    "\n",
    "dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)\n",
    "\n",
    "model_quantized_path = quantizer.quantize(\n",
    "    save_dir=save_dir+'_quantized',\n",
    "    quantization_config=dqconfig,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bddc2fae-e72a-4181-8120-cca22fd1db3c",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fc7f226c-9928-4f1a-bb8e-e52748fe9333",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using a model of type bert to instantiate a model of type . This is not supported for all configurations of models and can yield errors.\n",
      "The ONNX file model_optimized_quantized.onnx is not a regular name used in optimum.onnxruntime, the ORTModel might not behave as expected.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModel\n",
    "import torch\n",
    "\n",
    "\n",
    "#Mean Pooling - Take attention mask into account for correct averaging\n",
    "def mean_pooling(model_output, attention_mask):\n",
    "    token_embeddings = model_output[0] #First element of model_output contains all token embeddings\n",
    "    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()\n",
    "    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)\n",
    "\n",
    "\n",
    "# Sentences we want sentence embeddings for\n",
    "sentences = ['This is an example sentence', 'Each sentence is converted']\n",
    "\n",
    "# Load model from HuggingFace Hub\n",
    "tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-base-zh_optimized_quantized', model_max_length=512)\n",
    "\n",
    "model = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-base-zh_optimized_quantized')\n",
    "\n",
    "\n",
    "def generate_embeddings(tokenizer, model, sentences):\n",
    "\n",
    "    # Tokenize sentences\n",
    "    encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')\n",
    "    \n",
    "    # Compute token embeddings\n",
    "    with torch.no_grad():\n",
    "        model_output = model(**encoded_input)\n",
    "    \n",
    "    # Perform pooling. In this case, mean pooling.\n",
    "    sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])\n",
    "\n",
    "    return sentence_embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b8737642-1683-4400-b904-090ed6abec94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence embeddings:\n",
      "tensor([[ 0.2973,  0.0801,  0.6380,  ..., -0.7224, -0.6210, -0.3201],\n",
      "        [-0.1563,  0.2827,  0.5223,  ..., -0.8310, -0.6913, -0.2396]])\n",
      "768\n",
      "768\n"
     ]
    }
   ],
   "source": [
    "print(\"Sentence embeddings:\")\n",
    "sentence_embeddings = generate_embeddings(tokenizer, model, sentences)\n",
    "print(sentence_embeddings)\n",
    "print(len(sentence_embeddings[0]))\n",
    "print(len(sentence_embeddings[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe69fd91-a50e-4d26-8e94-0d68c6f4cd4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U sentence-transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "02108ba4-0d44-4222-8089-872ff716a7e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a8afc76220644434a0207b33ecebcd39",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "modules.json:   0%|          | 0.00/349 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d2f6753e1094a3d9800fb03d38acf27",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config_sentence_transformers.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d530521c84b448c48ddeaf19a73d4eda",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/27.9k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb39a285820347e18a3cfc9ac69ecf9b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "sentence_bert_config.json:   0%|          | 0.00/52.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e35b1b7591b45bfabca7b4768a4407c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "1_Pooling/config.json:   0%|          | 0.00/190 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "unoptimized_model = SentenceTransformer(model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "73703e86-ac15-4ace-92e0-1ad6836ac83e",
   "metadata": {},
   "outputs": [],
   "source": [
    "unopt_sent_emb = unoptimized_model.encode(sentences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e300911a-9457-42c1-90b4-39a3604a263a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics.pairwise import cosine_similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b4a53aec-5534-4d3a-b3c2-448d7e94aa07",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.745256781578064"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cosine_similarity([sentence_embeddings[1]],[unopt_sent_emb[1]]).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9979d8b7-5d9c-4d48-93c7-94b5a48a70af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}