Added an onnx model

.gitattributes

@@ -14,6 +14,7 @@
 *.npy filter=lfs diff=lfs merge=lfs -text
 *.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
+*.onnx.data filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
@@ -32,4 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+onnx/frida-onnx/FRIDA.onnx.data filter=lfs diff=lfs merge=lfs -text
+onnx/frida-onnx/FRIDA.onnx filter=lfs diff=lfs merge=lfs -text
+img.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,13 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+.idea
+uv.lock
+onnx_to_trt.py

README.md

@@ -1,3 +1,65 @@
----
-license: mit
----
+---
+license: mit
+language:
+- ru
+- en
+tags:
+- mteb
+- transformers
+- sentence-transformers
+base_model:
+- ai-forever/FRIDA
+pipeline_tag: feature-extraction
+---
+
+# FRIDA transformed to onnx
+Link to an [original](https://huggingface.co/ai-forever/FRIDA) repository for this model.
+This onnx version has batching support
+
+### Transform FRIDA to onnx and tensorrt
+This is a repository that contains FRIDA model in onnx (tensorrt upcoming) format.
+Python transformation scripts for this model are here too. onnx_to_trt.py is untested as of now
+
+# Model Card for FRIDA ONNX/TRT
+
+<figure>
+  <img src="img.jpg">
+</figure>
+
+FRIDA is a full-scale finetuned general text embedding model inspired by denoising architecture based on T5. The model is based on the encoder part of [FRED-T5](https://arxiv.org/abs/2309.10931) model and continues research of text embedding models ([ruMTEB](https://arxiv.org/abs/2408.12503), [ru-en-RoSBERTa](https://huggingface.co/ai-forever/ru-en-RoSBERTa)). It has been pre-trained on a Russian-English dataset and fine-tuned for improved performance on the target task. 
+
+For more model details please refer to this [article](https://habr.com/ru/companies/sberdevices/articles/909924/) (RU).
+
+## Usage
+
+The model can be used as is with prefixes. It is recommended to use CLS pooling. The choice of prefix and pooling depends on the task.
+
+We use the following basic rules to choose a prefix:
+- `"search_query: "` and `"search_document: "` prefixes are for answer or relevant paragraph retrieval
+- `"paraphrase: "` prefix is for symmetric paraphrasing related tasks (STS, paraphrase mining, deduplication)
+- `"categorize: "` prefix is for asymmetric matching of document title and body (e.g. news, scientific papers, social posts)
+- `"categorize_sentiment: "` prefix is for any tasks that rely on sentiment features (e.g. hate, toxic, emotion)
+- `"categorize_topic: "` prefix is intended for tasks where you need to group texts by topic
+- `"categorize_entailment: "` prefix is for textual entailment task (NLI)
+
+To better tailor the model to your needs, you can fine-tune it with relevant high-quality Russian and English datasets.
+
+Below are examples of texts encoding using the Transformers and SentenceTransformers libraries.
+
+## Authors
++ [SaluteDevices](https://sberdevices.ru/) AI for B2C RnD Team.
++ Artem Snegirev: [HF profile](https://huggingface.co/artemsnegirev), [Github](https://github.com/artemsnegirev);
++ Anna Maksimova [HF profile](https://huggingface.co/anpalmak);
++ Aleksandr Abramov: [HF profile](https://huggingface.co/Andrilko), [Github](https://github.com/Ab1992ao), [Kaggle Competitions Master](https://www.kaggle.com/andrilko)
+
+
+## Citation
+
+```
+@misc{TODO
+}
+```
+
+## Limitations
+
+The model is designed to process texts in Russian, the quality in English is unknown. Maximum input text length is limited to 512 tokens.

check_input_dims.py

@@ -0,0 +1,17 @@
+import onnxruntime as ort
+
+session = ort.InferenceSession("onnx/frida-onnx/FRIDA.onnx")
+
+for i, inp in enumerate(session.get_inputs()):
+    print(f"Input {i}:")
+    print(f"  Name: {inp.name}")
+    print(f"  Type: {inp.type}")
+    print(f"  Shape: {inp.shape}")
+    print(f"  Is dynamic? {'Yes' if -1 in inp.shape else 'No'}")
+
+for i, out in enumerate(session.get_outputs()):
+    print(f"Output {i}:")
+    print(f"  Name: {out.name}")
+    print(f"  Type: {out.type}")
+    print(f"  Shape: {out.shape}")
+    print(f"  Is dynamic? {'Yes' if -1 in out.shape else 'No'}")

check_inputs.py

@@ -0,0 +1,6 @@
+import onnx
+
+model = onnx.load("onnx/frida-onnx/FRIDA.onnx")
+for inp in model.graph.input:
+    shape = [dim.dim_value if dim.dim_value > 0 else str(dim.dim_param) for dim in inp.type.tensor_type.shape.dim]
+    print(f"Input '{inp.name}': {shape}")

onnx/frida-onnx/FRIDA.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aeaf461a810600c0489632bb3582b384bc531121cc3d7866a12ce0bcd6429461
+size 2514640

onnx/frida-onnx/FRIDA.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c01127d7452d91a0a5e596d292db84205f9c5672ac8ca43a5bab28f72b712e7
+size 3293628416

onnx/frida-onnx/special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

onnx/frida-onnx/tokenizer_config.json

@@ -0,0 +1,59 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<mask>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "mask_token": "<mask>",
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "repo_type": "model",
+  "sep_token": "</s>",
+  "tokenizer_class": "RobertaTokenizer",
+  "trim_offsets": true,
+  "unk_token": "<unk>"
+}

pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "frida-transformed"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13, <3.14"
+dependencies = [
+    'onnx == 1.20.0',
+    'onnxruntime == 1.23.2',
+    'onnxscript == 0.5.7',
+    'onnx-safetensors == 1.2.0',
+    'torch == 2.9.1',
+    'torchvision == 0.24.1',
+    'transformers == 4.57.3',
+    'tensorrt == 10.14.1.48.post1',
+    'pycuda == 2025.1.2'
+]
+
+[tool.uv.workspace]
+members = [
+    "frida-transformed",
+]

safetensors_to_onnx.py

@@ -0,0 +1,136 @@
+import torch
+from torch.export import Dim
+from transformers import T5EncoderModel, AutoTokenizer
+from pathlib import Path
+import onnxruntime as ort
+import numpy as np
+
+
+# MODEL_SOURCE_ID = "ai-forever/FRIDA"
+MODEL_SOURCE_ID = "../FRIDA"
+MODEL_TARGET_PATH = Path("onnx/frida-onnx")
+ONNX_FILE_NAME = "FRIDA.onnx"
+
+print("="*50)
+print(f"Подготовка директории: {MODEL_TARGET_PATH}")
+MODEL_TARGET_PATH.mkdir(parents=True, exist_ok=True)
+
+# 1. Загружаем модель и токенизатор
+print(f"Загрузка модели и токенизатора из '{MODEL_SOURCE_ID}'...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_SOURCE_ID, repo_type="model")
+model = T5EncoderModel.from_pretrained(MODEL_SOURCE_ID)
+model.eval()
+
+# 2. Создаем тестовые входы
+print("Создание тестовых входных данных...")
+test_texts = [
+    "paraphrase: В Ярославской области разрешили работу бань, но без посетителей",
+    "search_query: Сколько программистов нужно, чтобы вкрутить лампочку?",
+    "categorize_entailment: Женщину доставили в больницу, за ее жизнь сейчас борются врачи."
+]
+
+dummy_inputs = tokenizer(
+    test_texts,
+    max_length=512,
+    padding="max_length",
+    truncation=True,
+    return_tensors="pt"
+)
+
+# 3. Экспорт с двумя входами
+onnx_model_path = MODEL_TARGET_PATH / ONNX_FILE_NAME
+print(f"Экспорт модели в ONNX формат: {onnx_model_path}")
+
+# For dynamic_shapes
+batch_size = Dim("batch_size", min=1, max=64)  # Optional: add min/max constraints
+sequence_length = Dim("sequence_length", min=2, max=512)
+
+# dynamic_shapes = {
+#     "input_ids": {0: batch_size, 1: sequence_length},
+#     "attention_mask": {0: batch_size, 1: sequence_length},
+#     "last_hidden_state": {0: batch_size, 1: sequence_length}
+# }
+
+# In case of issues use dynamo_export instead of dynamo=True
+torch.onnx.export(
+    model,
+    (dummy_inputs["input_ids"], dummy_inputs["attention_mask"]),
+    onnx_model_path.as_posix(),
+    input_names=["input_ids", "attention_mask"],
+    output_names=["last_hidden_state"],
+    opset_version=20, # Maybe update
+        dynamic_shapes = {
+        "input_ids": {0: batch_size, 1: sequence_length},
+        "attention_mask": {0: batch_size, 1: sequence_length}
+    },
+    verbose=False,
+    dynamo=True
+)
+
+# 4. Сохраняем токенизатор
+print(f"Сохранение токенизатора в '{MODEL_TARGET_PATH}'...")
+tokenizer.save_pretrained(MODEL_TARGET_PATH)
+
+print("Конвертация завершена успешно!")
+
+# 5. Тестирование и сравнение результатов
+print("\n" + "="*50)
+print("ТЕСТИРОВАНИЕ РЕЗУЛЬТАТОВ")
+
+def cls_pooling(hidden_state, attention_mask):
+    """CLS pooling для получения эмбеддингов"""
+    return hidden_state[:, 0]
+
+def normalize_embeddings(embeddings):
+    """Нормализация эмбеддингов"""
+    return embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
+
+# Тест с оригинальной моделью
+print("Тестирование оригинальной модели...")
+with torch.no_grad():
+    original_inputs = tokenizer(
+        test_texts,
+        max_length=512,
+        padding=True,
+        truncation=True,
+        return_tensors="pt"
+    )
+    original_outputs = model(**original_inputs)
+    original_embeddings = cls_pooling(
+        original_outputs.last_hidden_state,
+        original_inputs["attention_mask"]
+    )
+    original_embeddings = torch.nn.functional.normalize(original_embeddings, p=2, dim=1)
+
+# Тест с ONNX моделью
+print("Тестирование ONNX модели...")
+onnx_session = ort.InferenceSession(onnx_model_path.as_posix())
+
+onnx_inputs = tokenizer(
+    test_texts,
+    max_length=512,
+    padding=True,
+    truncation=True,
+    return_tensors="np"
+)
+
+
+onnx_inputs_int64 = {
+    "input_ids": onnx_inputs["input_ids"].astype(np.int64),
+    "attention_mask": onnx_inputs["attention_mask"].astype(np.int64)
+}
+
+onnx_outputs = onnx_session.run(None, onnx_inputs_int64)[0]
+
+onnx_embeddings = onnx_outputs[:, 0]
+onnx_embeddings = normalize_embeddings(onnx_embeddings)
+
+cosine_similarity = np.sum(original_embeddings.numpy() * onnx_embeddings, axis=1)
+print(f"\nCosine similarity между оригинальной и ONNX моделью:")
+for i, sim in enumerate(cosine_similarity):
+    print(f"  Текст {i+1}: {sim:.6f}")
+print(f"Средняя схожесть: {np.mean(cosine_similarity):.6f}")
+
+print("\n" + "="*50)
+print("ГОТОВО! Модель успешно конвертирована и протестирована.")
+print(f"Путь к модели: {MODEL_TARGET_PATH.resolve()}")