onnx_convert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\imars\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer,AutoModelForTokenClassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "data = pd.read_csv('data_with_entity.csv')\n",
    "\n",
    "v_data = [' '.join(data['sent'].sample().to_list()) for i in range(100)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_from_torch_to_onnx(\n",
    "        onnx_path: str,\n",
    "        tokenizer: AutoTokenizer,\n",
    "        model: AutoModelForTokenClassification,\n",
    "        text: list\n",
    ") -> None:\n",
    "    \"\"\"Конвертация модели из формата PyTorch в формат ONNX.\n",
    "\n",
    "    @param onnx_path: путь к модели в формате ONNX\n",
    "    @param tokenizer: токенизатор\n",
    "    @param model: модель\n",
    "    \"\"\"\n",
    "    dummy_model_input = tokenizer(\n",
    "        text,\n",
    "        padding=\"max_length\",\n",
    "        truncation=True,\n",
    "        max_length=512,\n",
    "        return_tensors=\"pt\",\n",
    "    ).to(\"cpu\")\n",
    "    torch.onnx.export(\n",
    "        model,\n",
    "        dummy_model_input[\"input_ids\"],\n",
    "        onnx_path,\n",
    "        opset_version=12,\n",
    "        input_names=[\"input_ids\"],\n",
    "        output_names=[\"output\"],\n",
    "        dynamic_axes={\n",
    "            \"input_ids\": {\n",
    "                0: \"batch_size\",\n",
    "                1: \"sequence_len\"\n",
    "            },\n",
    "            \"output\": {\n",
    "                0: \"batch_size\"\n",
    "            }\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\imars\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\huggingface_hub\\file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n",
      "Some weights of ElectraForTokenClassification were not initialized from the model checkpoint at ai-forever/ruElectra-small and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model_name = 'ai-forever/ruElectra-small'\n",
    "label2idx = {'O': 0, 'B-discount': 1, 'B-value': 2, 'I-value': 3}\n",
    "model = AutoModelForTokenClassification.from_pretrained(model_name, num_labels = len(label2idx)).to('cpu')\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert_from_torch_to_onnx('ruElectra-small-onnx.onnx', tokenizer, model, v_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from onnxruntime.quantization import (\n",
    "    quantize_dynamic,\n",
    "    QuantType\n",
    ")\n",
    "\n",
    "\n",
    "def convert_from_onnx_to_quantized_onnx(\n",
    "        onnx_model_path: str,\n",
    "        quantized_onnx_model_path: str\n",
    ") -> None:\n",
    "    \"\"\"Квантизация модели в формате ONNX до Int8\n",
    "    и сохранение кванитизированной модели на диск.\n",
    "\n",
    "    @param onnx_model_path: путь к модели в формате ONNX\n",
    "    @param quantized_onnx_model_path: путь к квантизированной модели\n",
    "    \"\"\"\n",
    "    quantize_dynamic(\n",
    "        onnx_model_path,\n",
    "        quantized_onnx_model_path,\n",
    "        weight_type=QuantType.QUInt8\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Please consider to run pre-processing before quantization. Refer to example: https://github.com/microsoft/onnxruntime-inference-examples/blob/main/quantization/image_classification/cpu/ReadMe.md \n"
     ]
    }
   ],
   "source": [
    "convert_from_onnx_to_quantized_onnx('ruElectra-small-onnx.onnx', 'ruElectra-small-onnx-quantized.onnx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoModelForSequenceClassification\n",
    ")\n",
    "\n",
    "\n",
    "torch.set_num_threads(1)\n",
    "\n",
    "\n",
    "def pytorch_inference(\n",
    "        text: str,\n",
    "        max_tokens: int,\n",
    "        model: AutoModelForSequenceClassification,\n",
    "        tokenizer: AutoTokenizer,\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"Инференс модели с помощью PyTorch.\n",
    "\n",
    "    @param text: входной текст для классификации\n",
    "    @param max_tokens: максимальная длина последовательности в токенах\n",
    "    @param model: BERT-модель\n",
    "    @param tokenizer: токенизатор\n",
    "    @return: логиты на выходе из модели\n",
    "    \"\"\"\n",
    "    inputs = tokenizer(\n",
    "        text,\n",
    "        padding=\"max_length\",\n",
    "        truncation=True,\n",
    "        max_length=max_tokens,\n",
    "        return_tensors=\"pt\"\n",
    "    ).to('cpu')\n",
    "    with torch.inference_mode():\n",
    "        outputs = model(**inputs).logits.detach()\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime\n",
    "from onnxruntime import (\n",
    "    InferenceSession,\n",
    "    SessionOptions\n",
    ")\n",
    "\n",
    "\n",
    "def create_onnx_session(\n",
    "        model_path: str,\n",
    "        provider: str = \"CPUExecutionProvider\"\n",
    ") -> InferenceSession:\n",
    "    \"\"\"Создание сессии для инференса модели с помощью ONNX Runtime.\n",
    "\n",
    "    @param model_path: путь к модели в формате ONNX\n",
    "    @param provider: инференс на ЦП\n",
    "    @return: ONNX Runtime-сессия\n",
    "    \"\"\"    \n",
    "    options = SessionOptions()\n",
    "    options.graph_optimization_level = \\\n",
    "        onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL\n",
    "    options.intra_op_num_threads = 1\n",
    "    session = InferenceSession(model_path, options, providers=[provider])\n",
    "    session.disable_fallback()\n",
    "    return session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from transformers import AutoTokenizer\n",
    "from onnxruntime import InferenceSession\n",
    "\n",
    "\n",
    "def onnx_inference(\n",
    "        text: str,\n",
    "        session: InferenceSession,\n",
    "        tokenizer: AutoTokenizer,\n",
    "        max_length: int\n",
    ") -> np.ndarray:\n",
    "    \"\"\"Инференс модели с помощью ONNX Runtime.\n",
    "\n",
    "    @param text: входной текст для классификации\n",
    "    @param session: ONNX Runtime-сессия\n",
    "    @param tokenizer: токенизатор\n",
    "    @param max_length: максимальная длина последовательности в токенах\n",
    "    @return: логиты на выходе из модели\n",
    "    \"\"\"\n",
    "    inputs = tokenizer(\n",
    "        text,\n",
    "        padding=\"max_length\",\n",
    "        truncation=True,\n",
    "        max_length=max_length,\n",
    "        return_tensors=\"np\",\n",
    "    )\n",
    "    input_feed = {\n",
    "        \"input_ids\": inputs[\"input_ids\"].astype(np.int64)\n",
    "    }\n",
    "    outputs = session.run(\n",
    "        output_names=[\"output\"],\n",
    "        input_feed=input_feed\n",
    "    )[0]\n",
    "    return outputs"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}