Monor improvement

safetensors_to_onnx.ipynb

@@ -2,11 +2,7 @@
  "cells": [
   {
    "metadata": {
-    "collapsed": true,
-    "ExecuteTime": {
-     "end_time": "2026-02-12T12:52:46.678786554Z",
-     "start_time": "2026-02-12T12:52:43.490350354Z"
-    }
+    "collapsed": true
    },
    "cell_type": "code",
    "source": [
@@ -20,15 +16,10 @@
    ],
    "id": "2b3977272abf14d9",
    "outputs": [],
-   "execution_count": 1
+   "execution_count": null
   },
   {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T12:52:46.726391124Z",
-     "start_time": "2026-02-12T12:52:46.691717774Z"
-    }
-   },
+   "metadata": {},
    "cell_type": "code",
    "source": [
     "# MODEL_SOURCE_ID = \"sergeyzh/BERTA\"\n",
@@ -41,25 +32,11 @@
     "MODEL_TARGET_PATH.mkdir(parents=True, exist_ok=True)"
    ],
    "id": "494fc15203b0fb89",
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "==================================================\n",
-      "Подготовка директории: onnx/berta-onnx\n"
-     ]
-    }
-   ],
-   "execution_count": 2
+   "outputs": [],
+   "execution_count": null
   },
   {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T12:52:46.862603179Z",
-     "start_time": "2026-02-12T12:52:46.739714466Z"
-    }
-   },
+   "metadata": {},
    "cell_type": "code",
    "source": [
     "# 1. Загружаем модель и токенизатор\n",
@@ -86,56 +63,22 @@
     "print(dummy_inputs)"
    ],
    "id": "4f9f5febc6f07769",
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Загрузка модели и токенизатора из '../BERTA'...\n",
-      "Создание тестовых входных данных...\n",
-      "{'input_ids': tensor([[    2,   570, 11028,  ...,     0,     0,     0],\n",
-      "        [    2,  3007,    67,  ...,     0,     0,     0],\n",
-      "        [    2, 46369,   998,  ...,     0,     0,     0]]), 'token_type_ids': tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
-      "        [0, 0, 0,  ..., 0, 0, 0],\n",
-      "        [0, 0, 0,  ..., 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
-      "        [1, 1, 1,  ..., 0, 0, 0],\n",
-      "        [1, 1, 1,  ..., 0, 0, 0]])}\n"
-     ]
-    }
-   ],
-   "execution_count": 3
+   "outputs": [],
+   "execution_count": null
   },
   {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T12:52:46.899958136Z",
-     "start_time": "2026-02-12T12:52:46.868506089Z"
-    }
-   },
+   "metadata": {},
    "cell_type": "code",
    "source": [
     "# print(model)\n",
     "print(signature(model.forward))"
    ],
    "id": "8bdce4e5bc593383",
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[transformers.cache_utils.Cache] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.Tensor] = None) -> Union[tuple[torch.Tensor], transformers.modeling_outputs.BaseModelOutputWithPoolingAndCrossAttentions]\n"
-     ]
-    }
-   ],
-   "execution_count": 4
+   "outputs": [],
+   "execution_count": null
   },
   {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T12:52:56.427369911Z",
-     "start_time": "2026-02-12T12:52:46.902043777Z"
-    }
-   },
+   "metadata": {},
    "cell_type": "code",
    "source": [
     "# 3. Экспорт с двумя входами\n",
@@ -153,20 +96,21 @@
     "# }\n",
     "\n",
     "# In case of issues use dynamo_export instead of dynamo=True\n",
-    "torch.onnx.export(\n",
-    "    model,\n",
-    "    (dummy_inputs[\"input_ids\"], dummy_inputs[\"attention_mask\"]),\n",
-    "    onnx_model_path.as_posix(),\n",
-    "    input_names=[\"input_ids\", \"attention_mask\"],\n",
-    "    output_names=[\"last_hidden_state\"],\n",
-    "    opset_version=20, # Maybe update\n",
-    "    dynamic_shapes = {\n",
-    "        \"input_ids\": {0: batch_size, 1: sequence_length},\n",
-    "        \"attention_mask\": {0: batch_size, 1: sequence_length}\n",
-    "    },\n",
-    "    verbose=True,\n",
-    "    dynamo=True\n",
-    ")\n",
+    "with torch.no_grad():\n",
+    "    torch.onnx.export(\n",
+    "        model,\n",
+    "        (dummy_inputs[\"input_ids\"], dummy_inputs[\"attention_mask\"]),\n",
+    "        onnx_model_path.as_posix(),\n",
+    "        input_names=[\"input_ids\", \"attention_mask\"],\n",
+    "        output_names=[\"last_hidden_state\"],\n",
+    "        opset_version=20, # Maybe update\n",
+    "        dynamic_shapes = {\n",
+    "            \"input_ids\": {0: batch_size, 1: sequence_length},\n",
+    "            \"attention_mask\": {0: batch_size, 1: sequence_length}\n",
+    "        },\n",
+    "        verbose=True,\n",
+    "        dynamo=True\n",
+    "    )\n",
     "# 4. Сохраняем токенизатор\n",
     "print(f\"Сохранение токенизатора в '{MODEL_TARGET_PATH}'...\")\n",
     "tokenizer.save_pretrained(MODEL_TARGET_PATH)\n",
@@ -174,97 +118,11 @@
     "print(\"Конвертация завершена успешно!\")"
    ],
    "id": "87d59bf71ed545dc",
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Экспорт модели в ONNX формат: onnx/berta-onnx/BERTA.onnx\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "W0212 14:52:47.799000 19280 torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'input' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0, sampling_ratio: 'int' = -1, aligned: 'bool' = False). Treating as an Input.\n",
-      "W0212 14:52:47.800000 19280 torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'boxes' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0, sampling_ratio: 'int' = -1, aligned: 'bool' = False). Treating as an Input.\n",
-      "W0212 14:52:47.801000 19280 torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'input' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0). Treating as an Input.\n",
-      "W0212 14:52:47.802000 19280 torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'boxes' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0). Treating as an Input.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[torch.onnx] Obtain model graph for `BertModel([...]` with `torch.export.export(..., strict=False)`...\n",
-      "[torch.onnx] Obtain model graph for `BertModel([...]` with `torch.export.export(..., strict=False)`... ✅\n",
-      "[torch.onnx] Run decomposition...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/lavrentiy/Projects/FRIDA-transformed/.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:435: UserWarning: \n",
-      "    Found GPU0 NVIDIA GeForce GTX 1060 6GB which is of cuda capability 6.1.\n",
-      "    Minimum and Maximum cuda capability supported by this version of PyTorch is\n",
-      "    (7.0) - (12.0)\n",
-      "    \n",
-      "  queued_call()\n",
-      "/home/lavrentiy/Projects/FRIDA-transformed/.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:435: UserWarning: \n",
-      "    Please install PyTorch with a following CUDA\n",
-      "    configurations:  12.6 following instructions at\n",
-      "    https://pytorch.org/get-started/locally/\n",
-      "    \n",
-      "  queued_call()\n",
-      "/home/lavrentiy/Projects/FRIDA-transformed/.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:435: UserWarning: \n",
-      "NVIDIA GeForce GTX 1060 6GB with CUDA capability sm_61 is not compatible with the current PyTorch installation.\n",
-      "The current PyTorch install supports CUDA capabilities sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120.\n",
-      "If you want to use the NVIDIA GeForce GTX 1060 6GB GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/\n",
-      "\n",
-      "  queued_call()\n",
-      "/home/lavrentiy/.local/share/uv/python/cpython-3.13.11-linux-x86_64-gnu/lib/python3.13/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.\n",
-      "  return cls.__new__(cls, *args)\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[torch.onnx] Run decomposition... ✅\n",
-      "[torch.onnx] Translate the graph into ONNX...\n",
-      "[torch.onnx] Translate the graph into ONNX... ✅\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/lavrentiy/Projects/FRIDA-transformed/.venv/lib/python3.13/site-packages/torch/onnx/_internal/exporter/_onnx_program.py:460: UserWarning: # The axis name: batch_size will not be used, since it shares the same shape constraints with another axis: batch_size.\n",
-      "  rename_mapping = _dynamic_shapes.create_rename_mapping(\n",
-      "/home/lavrentiy/Projects/FRIDA-transformed/.venv/lib/python3.13/site-packages/torch/onnx/_internal/exporter/_onnx_program.py:460: UserWarning: # The axis name: sequence_length will not be used, since it shares the same shape constraints with another axis: sequence_length.\n",
-      "  rename_mapping = _dynamic_shapes.create_rename_mapping(\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Applied 68 of general pattern rewrite rules.\n",
-      "Сохранение токенизатора в 'onnx/berta-onnx'...\n",
-      "Конвертация завершена успешно!\n"
-     ]
-    }
-   ],
-   "execution_count": 5
+   "outputs": [],
+   "execution_count": null
   },
   {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T12:52:56.931194388Z",
-     "start_time": "2026-02-12T12:52:56.428745759Z"
-    }
-   },
+   "metadata": {},
    "cell_type": "code",
    "source": [
     "# 5. Тестирование и сравнение результатов\n",
@@ -330,30 +188,8 @@
     "print(f\"Путь к модели: {MODEL_TARGET_PATH.resolve()}\")"
    ],
    "id": "91a5740805f8e829",
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n",
-      "==================================================\n",
-      "ТЕСТИРОВАНИЕ РЕЗУЛЬТАТОВ\n",
-      "Тестирование оригинальной модели...\n",
-      "Тестирование ONNX модели...\n",
-      "\n",
-      "Cosine similarity между оригинальной и ONNX моделью:\n",
-      "  Текст 1: 1.000000\n",
-      "  Текст 2: 1.000000\n",
-      "  Текст 3: 1.000000\n",
-      "Средняя схожесть: 1.000000\n",
-      "\n",
-      "==================================================\n",
-      "ГОТОВО! Модель успешно конвертирована и протестирована.\n",
-      "Путь к модели: /home/lavrentiy/Projects/BERTA-transformed/onnx/berta-onnx\n"
-     ]
-    }
-   ],
-   "execution_count": 6
+   "outputs": [],
+   "execution_count": null
   }
  ],
  "metadata": {