Added signature for viewing model inputs

pyproject.toml

@@ -5,15 +5,21 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.13, <3.14"
 dependencies = [
-    'onnx == 1.20.0',
+    'onnx == 1.20.1',
     'onnxruntime == 1.23.2',
-    'onnxscript == 0.5.7',
-    'onnx-safetensors == 1.2.0',
-    'torch == 2.9.1',
-    'torchvision == 0.24.1',
+    'onnxscript == 0.6.0',
+    'onnx-safetensors == 1.5.0',
+    'torch == 2.10.0',
+    'torchvision == 0.25.0',
     'transformers == 4.57.3',
-    'tensorrt == 10.14.1.48.post1',
-    'pycuda == 2025.1.2'
+    'pycuda == 2026.1',
+    "ipykernel>=7.2.0",
+    "pip>=26.0.1",
+    "uv>=0.10.2",
+    "jupyter>=1.1.1",
+    "ipywidgets>=8.1.8",
+    "tqdm>=4.67.3",
+    "ipython>=9.10.0",
 ]
 
 [tool.uv.workspace]

safetensors_to_onnx.ipynb

@@ -0,0 +1,257 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true,
+    "ExecuteTime": {
+     "end_time": "2026-02-12T12:37:32.166521648Z",
+     "start_time": "2026-02-12T12:37:32.138056109Z"
+    }
+   },
+   "source": [
+    "import torch\n",
+    "from torch.export import Dim\n",
+    "from transformers import T5EncoderModel, AutoTokenizer\n",
+    "from pathlib import Path\n",
+    "import onnxruntime as ort\n",
+    "import numpy as np\n",
+    "from inspect import signature"
+   ],
+   "outputs": [],
+   "execution_count": 5
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-02-12T12:37:00.482648074Z",
+     "start_time": "2026-02-12T12:37:00.118707317Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# MODEL_SOURCE_ID = \"ai-forever/FRIDA\"\n",
+    "MODEL_SOURCE_ID = \"../FRIDA\"\n",
+    "MODEL_TARGET_PATH = Path(\"onnx/frida-onnx\")\n",
+    "ONNX_FILE_NAME = \"FRIDA.onnx\"\n",
+    "\n",
+    "print(\"=\"*50)\n",
+    "print(f\"Подготовка директории: {MODEL_TARGET_PATH}\")\n",
+    "MODEL_TARGET_PATH.mkdir(parents=True, exist_ok=True)"
+   ],
+   "id": "ef5e190f02e042b6",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "==================================================\n",
+      "Подготовка директории: onnx/frida-onnx\n"
+     ]
+    }
+   ],
+   "execution_count": 2
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-02-12T12:37:17.778488452Z",
+     "start_time": "2026-02-12T12:37:16.890360137Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# 1. Загружаем модель и токенизатор\n",
+    "print(f\"Загрузка модели и токенизатора из '{MODEL_SOURCE_ID}'...\")\n",
+    "tokenizer = AutoTokenizer.from_pretrained(MODEL_SOURCE_ID, repo_type=\"model\")\n",
+    "model = T5EncoderModel.from_pretrained(MODEL_SOURCE_ID)\n",
+    "model.eval()\n",
+    "\n",
+    "# 2. Создаем тестовые входы\n",
+    "print(\"Создание тестовых входных данных...\")\n",
+    "test_texts = [\n",
+    "    \"paraphrase: В Ярославской области разрешили работу бань, но без посетителей\",\n",
+    "    \"search_query: Сколько программистов нужно, чтобы вкрутить лампочку?\",\n",
+    "    \"categorize_entailment: Женщину доставили в больницу, за ее жизнь сейчас борются врачи.\"\n",
+    "]\n",
+    "\n",
+    "dummy_inputs = tokenizer(\n",
+    "    test_texts,\n",
+    "    max_length=512,\n",
+    "    padding=\"max_length\",\n",
+    "    truncation=True,\n",
+    "    return_tensors=\"pt\"\n",
+    ")"
+   ],
+   "id": "d2913ab82e279832",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Загрузка модели и токенизатора из '../FRIDA'...\n",
+      "Создание тестовых входных данных...\n"
+     ]
+    }
+   ],
+   "execution_count": 3
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-02-12T12:37:34.830442932Z",
+     "start_time": "2026-02-12T12:37:34.719042026Z"
+    }
+   },
+   "cell_type": "code",
+   "source": "print(signature(model.forward))",
+   "id": "e55cf99269a639d2",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None) -> Union[tuple[torch.FloatTensor], transformers.modeling_outputs.BaseModelOutput]\n"
+     ]
+    }
+   ],
+   "execution_count": 6
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# 3. Экспорт с двумя входами\n",
+    "onnx_model_path = MODEL_TARGET_PATH / ONNX_FILE_NAME\n",
+    "print(f\"Экспорт модели в ONNX формат: {onnx_model_path}\")\n",
+    "\n",
+    "# For dynamic_shapes\n",
+    "batch_size = Dim(\"batch_size\", min=1, max=64)  # Optional: add min/max constraints\n",
+    "sequence_length = Dim(\"sequence_length\", min=2, max=512)\n",
+    "\n",
+    "# dynamic_shapes = {\n",
+    "#     \"input_ids\": {0: batch_size, 1: sequence_length},\n",
+    "#     \"attention_mask\": {0: batch_size, 1: sequence_length},\n",
+    "#     \"last_hidden_state\": {0: batch_size, 1: sequence_length}\n",
+    "# }\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=False,\n",
+    "    dynamo=True\n",
+    ")\n",
+    "\n",
+    "# 4. Сохраняем токенизатор\n",
+    "print(f\"Сохранение токенизатора в '{MODEL_TARGET_PATH}'...\")\n",
+    "tokenizer.save_pretrained(MODEL_TARGET_PATH)\n",
+    "\n",
+    "print(\"Конвертация завершена успешно!\")"
+   ],
+   "id": "48bfef4b286ae47b",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# 5. Тестирование и сравнение результатов\n",
+    "print(\"\\n\" + \"=\"*50)\n",
+    "print(\"ТЕСТИРОВАНИЕ РЕЗУЛЬТАТОВ\")\n",
+    "\n",
+    "def cls_pooling(hidden_state, attention_mask):\n",
+    "    \"\"\"CLS pooling для получения эмбеддингов\"\"\"\n",
+    "    return hidden_state[:, 0]\n",
+    "\n",
+    "def normalize_embeddings(embeddings):\n",
+    "    \"\"\"Нормализация эмбеддингов\"\"\"\n",
+    "    return embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)\n",
+    "\n",
+    "# Тест с оригинальной моделью\n",
+    "print(\"Тестирование оригинальной модели...\")\n",
+    "with torch.no_grad():\n",
+    "    original_inputs = tokenizer(\n",
+    "        test_texts,\n",
+    "        max_length=512,\n",
+    "        padding=True,\n",
+    "        truncation=True,\n",
+    "        return_tensors=\"pt\"\n",
+    "    )\n",
+    "    original_outputs = model(**original_inputs)\n",
+    "    original_embeddings = cls_pooling(\n",
+    "        original_outputs.last_hidden_state,\n",
+    "        original_inputs[\"attention_mask\"]\n",
+    "    )\n",
+    "    original_embeddings = torch.nn.functional.normalize(original_embeddings, p=2, dim=1)\n",
+    "\n",
+    "# Тест с ONNX моделью\n",
+    "print(\"Тестирование ONNX модели...\")\n",
+    "onnx_session = ort.InferenceSession(onnx_model_path.as_posix())\n",
+    "\n",
+    "onnx_inputs = tokenizer(\n",
+    "    test_texts,\n",
+    "    max_length=512,\n",
+    "    padding=True,\n",
+    "    truncation=True,\n",
+    "    return_tensors=\"np\"\n",
+    ")\n",
+    "\n",
+    "\n",
+    "onnx_inputs_int64 = {\n",
+    "    \"input_ids\": onnx_inputs[\"input_ids\"].astype(np.int64),\n",
+    "    \"attention_mask\": onnx_inputs[\"attention_mask\"].astype(np.int64)\n",
+    "}\n",
+    "\n",
+    "onnx_outputs = onnx_session.run(None, onnx_inputs_int64)[0]\n",
+    "\n",
+    "onnx_embeddings = onnx_outputs[:, 0]\n",
+    "onnx_embeddings = normalize_embeddings(onnx_embeddings)\n",
+    "\n",
+    "cosine_similarity = np.sum(original_embeddings.numpy() * onnx_embeddings, axis=1)\n",
+    "print(f\"\\nCosine similarity между оригинальной и ONNX моделью:\")\n",
+    "for i, sim in enumerate(cosine_similarity):\n",
+    "    print(f\"  Текст {i+1}: {sim:.6f}\")\n",
+    "print(f\"Средняя схожесть: {np.mean(cosine_similarity):.6f}\")\n",
+    "\n",
+    "print(\"\\n\" + \"=\"*50)\n",
+    "print(\"ГОТОВО! Модель успешно конвертирована и протестирована.\")\n",
+    "print(f\"Путь к модели: {MODEL_TARGET_PATH.resolve()}\")"
+   ],
+   "id": "e488535f18210818",
+   "outputs": [],
+   "execution_count": null
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}