Repository: unum-cloud/uform (language: python)

Issue:
Add: register tokens in ViT



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/prerelease.yml b/.github/workflows/prerelease.yml
new file mode 100644
index 0000000..9646d21
--- /dev/null
+++ b/.github/workflows/prerelease.yml
@@ -0,0 +1,38 @@
+name: Pre-Release
+
+on:
+  push:
+    branches: ["main-dev"]
+  pull_request:
+    branches: ["main-dev"]
+
+env:
+  BUILD_TYPE: Release
+  GH_TOKEN: ${{ secrets.SEMANTIC_RELEASE_TOKEN }}
+  PYTHONUTF8: 1
+
+# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
+permissions:
+  contents: read
+
+jobs:
+  test_python:
+    name: Test Python
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - run: git submodule update --init --recursive
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --no-cache-dir --upgrade pip
+          pip install -e ".[onnx]"
+          pip install pytest
+
+      - name: Test with PyTest
+        run: pytest python/scripts/ -s -x -Wd -v
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 230b303..df18e4b 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -1,11 +1,3 @@
-# This workflow will upload a Python Package using Twine when a release is created
-# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries
-
-# This workflow uses actions that are not certified by GitHub.
-# They are provided by a third-party and are governed by
-# separate terms of service, privacy policy, and support
-# documentation.
-
 name: Release
 
 on:
@@ -26,11 +18,13 @@ jobs:
     name: Semantic Release
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           persist-credentials: false
       - uses: actions/setup-node@v3
-      - run: cp .github/workflows/package.json . && npm install && npx semantic-release
+        with:
+          node-version: 20
+      - run: npm install --save-dev --prefix ./package-ci @semantic-release/exec @semantic-release/git conventional-changelog-eslint semantic-release && npx --prefix ./package-ci semantic-release
 
   rebase:
     name: Rebase Dev. Branch
@@ -38,7 +32,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout the latest code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
@@ -56,17 +50,40 @@ jobs:
           unprotect_reviews: True
           force: True
 
+  test_python:
+    name: Run Tests
+    runs-on: ubuntu-latest
+    needs: versioning
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: "main"
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[onnx]"
+          pip install pytest
+
+      - name: Run PyTest
+        run: pytest python/scripts/
+
   pypi_publish:
     name: Publish Python
     runs-on: ubuntu-latest
     needs: versioning
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           ref: "main"
       - name: Set up Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v5
         with:
           python-version: "3.11"
 
@@ -79,7 +96,7 @@ jobs:
         run: python -m build
 
       - name: Publish to PyPi
-        uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+        uses: pypa/gh-action-pypi-publish@1.8.14
         with:
           user: __token__
           password: ${{ secrets.PYPI_API_TOKEN }}
@@ -93,7 +110,7 @@ jobs:
     needs: pypi_publish
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           ref: "main"
       - name: Setup GitHub Pages
diff --git a/.github/workflows/update_version.sh b/.github/workflows/update_version.sh
new file mode 100755
index 0000000..ae04d66
--- /dev/null
+++ b/.github/workflows/update_version.sh
@@ -0,0 +1,5 @@
+#!/bin/sh
+
+echo $1 > VERSION && 
+    sed -i "s/^version: .*/version: $1/" CITATION.cff &&
+    sed -i "s/version = \".*\"/version = \"$1\"/" pyproject.toml
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 8e9fde4..1732614 100755
--- a/.gitignore
+++ b/.gitignore
@@ -1,8 +1,8 @@
 requirements.txt
-*.egg-info
 dist/
 test
-src/__pycache__
-src/test.py
 build/
 package-lock.json
+*.egg-info
+*.onnx
+__pycache__
\ No newline at end of file
diff --git a/.releaserc b/.releaserc
new file mode 100644
index 0000000..9b1f15e
--- /dev/null
+++ b/.releaserc
@@ -0,0 +1,84 @@
+{
+    "branches": [
+        "main"
+    ],
+    "debug": true,
+    "ci": true,
+    "dryRun": false,
+    "plugins": [
+        [
+            "@semantic-release/commit-analyzer",
+            {
+                "preset": "eslint",
+                "releaseRules": [
+                    {
+                        "tag": "Break",
+                        "release": "major"
+                    },
+                    {
+                        "tag": "Add",
+                        "release": "minor"
+                    },
+                    {
+                        "tag": "Improve",
+                        "release": "patch"
+                    },
+                    {
+                        "tag": "Make",
+                        "release": "patch"
+                    },
+                    {
+                        "tag": "Refactor",
+                        "release": false
+                    }
+                ]
+            }
+        ],
+        [
+            "@semantic-release/release-notes-generator",
+            {
+                "preset": "eslint",
+                "releaseRules": [
+                    {
+                        "tag": "Break",
+                        "release": "major"
+                    },
+                    {
+                        "tag": "Add",
+                        "release": "minor"
+                    },
+                    {
+                        "tag": "Improve",
+                        "release": "patch"
+                    },
+                    {
+                        "tag": "Make",
+                        "release": "patch"
+                    },
+                    {
+                        "tag": "Refactor",
+                        "release": false
+                    }
+                ]
+            }
+        ],
+        "@semantic-release/github",
+        [
+            "@semantic-release/exec",
+            {
+                "prepareCmd": "bash .github/workflows/update_version.sh '${nextRelease.version}'"
+            }
+        ],
+        [
+            "@semantic-release/git",
+            {
+                "assets": [
+                    "VERSION",
+                    "pyproject.toml",
+                    "CITATION.cff"
+                ],
+                "message": "Build: Released ${nextRelease.version} [skip ci]\n\n${nextRelease.notes}"
+            }
+        ]
+    ]
+}
diff --git a/.vscode/settings.json b/.vscode/settings.json
index d872cfc..9b08d04 100755
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,13 +1,23 @@
 {
     "cSpell.words": [
+        "dtype",
         "embs",
+        "huggingface",
         "keepdim",
         "logits",
+        "Matryoshka",
         "multimodal",
+        "ndarray",
+        "numpy",
+        "ONNX",
+        "onnxruntime",
         "preprocess",
         "pretrained",
         "probs",
         "pypi",
+        "rerank",
+        "reranker",
+        "reranking",
         "softmax",
         "transfromers",
         "uform",
@@ -18,4 +28,4 @@
         "editor.defaultFormatter": "ms-python.black-formatter"
     },
     "python.formatting.provider": "none"
-}
+}
\ No newline at end of file
diff --git a/CITATION.cff b/CITATION.cff
index 36dc871..ee6c162 100644
--- a/CITATION.cff
+++ b/CITATION.cff
@@ -11,7 +11,7 @@ authors:
   given-names: "Ash"
   orcid: "https://orcid.org/0000-0002-4882-1815"
 title: "UForm by Unum Cloud"
-version: 0.4.2
+version: 1.1.2
 keywords:
 - "text-to-image retrieval"
 - "multimodal"
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 0000000..1548c30
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,16 @@
+# Contributing to UForm
+
+We welcome contributions to UForm!
+Before submitting any changes, please make sure that the tests pass.
+
+```sh
+pip install -e .                # For core dependencies
+
+pip install -e ".[torch]"       # For PyTorch
+pip install -e ".[onnx]"        # For ONNX on CPU
+pip install -e ".[onnx-gpu]"    # For ONNX on GPU, available for some platforms
+pip install -e ".[torch,onnx]"  # For PyTorch and ONNX Python tests
+
+pytest python/scripts/ -s -x -Wd -v
+pytest python/scripts/ -s -x -Wd -v -k onnx # To run only ONNX tests without loading Torch
+```
\ No newline at end of file
diff --git a/README.md b/README.md
index 2d4dcf5..031c484 100755
--- a/README.md
+++ b/README.md
@@ -17,6 +17,14 @@ For Content Understanding and Generation<br/>
 <a href="https://github.com/unum-cloud/uform"><img height="25" src="https://github.com/unum-cloud/.github/raw/main/assets/github.svg" alt="GitHub"></a>
 </p>
 
+<p align="center">
+Multimodal Embeddings from 64 to 768 Dimensions • 1B Parameter Chat
+<br/>
+Short Texts • Images • 🔜 Video Clips
+<br/>
+PyTorch • ONNX
+</p>
+
 ---
 
 ![](https://github.com/ashvardanian/usearch-images/blob/main/assets/uform-gen-preview.jpg?raw=true)
@@ -28,14 +36,14 @@ With compact __custom pre-trained transformer models__, this can run anywhere fr
 
 ## Features
 
-* __Throughput__: Thanks to the small size, the inference speed is [2-4x faster](#speed) than competitors.
-* __Tiny Embeddings__: 256-dimensional vectors are 2-3x quicker to [search][usearch] than from CLIP-like models.
-* __Quantization Aware__: Downcasted embeddings from `f32` to `i8` without losing much recall.
-* __Multilingual__: Trained on a balanced dataset, the recall is great across over [20 languages](#evaluation).
-* __Hardware Friendly__: Whether it's Apple's CoreML or ONNX, [we've got you covered][onnx].
+- __Tiny Embeddings__: 64-dimensional [Matryoshaka][matryoshka]-style embeddings for extremely fast [search][usearch].
+- __Throughput__: Thanks to the small size, the inference speed is [2-4x faster](#speed) than competitors.
+- __Portable__: Models come with native ONNX support, making them easy to deploy on any platform.
+- __Quantization Aware__: Down-cast embeddings from `f32` to `i8` without losing much recall.
+- __Multilingual__: Trained on a balanced dataset, the recall is great across over [20 languages](#evaluation).
 
 [usearch]: https://github.com/unum-cloud/usearch
-[onnx]: https://huggingface.co/unum-cloud/uform-coreml-onnx
+[matryoshka]: https://arxiv.org/abs/2205.13147
 
 ## Models
 
@@ -43,83 +51,208 @@ With compact __custom pre-trained transformer models__, this can run anywhere fr
 
 | Model                                    | Parameters | Languages |                                 Architecture |
 | :--------------------------------------- | ---------: | --------: | -------------------------------------------: |
+| [`uform-vl-english-large`][model-e-l] 🆕  |       365M |         1 | 6 text layers, ViT-L/14, 6 multimodal layers |
 | [`uform-vl-english`][model-e]            |       143M |         1 | 2 text layers, ViT-B/16, 2 multimodal layers |
+| [`uform-vl-english-small`][model-e-s] 🆕  |        79M |         1 | 2 text layers, ViT-S/16, 2 multimodal layers |
 | [`uform-vl-multilingual-v2`][model-m-v2] |       206M |        21 | 8 text layers, ViT-B/16, 4 multimodal layers |
 | [`uform-vl-multilingual`][model-m]       |       206M |        12 | 8 text layers, ViT-B/16, 4 multimodal layers |
 
+[model-e-l]: https://huggingface.co/unum-cloud/uform-vl-english-large/
 [model-e]: https://huggingface.co/unum-cloud/uform-vl-english/
+[model-e-s]: https://huggingface.co/unum-cloud/uform-vl-english-small/
 [model-m]: https://huggingface.co/unum-cloud/uform-vl-multilingual/
 [model-m-v2]: https://huggingface.co/unum-cloud/uform-vl-multilingual-v2/
 
 ### Generative Models
 
-| Model                              | Parameters |            Purpose          |     Architecture      |
-| :--------------------------------- | ---------: | --------------------------: | --------------------: |
-| [`uform-gen2-qwen-500m`][model-g2] |    1.2B    | Chat, Image Captioning, VQA | qwen1.5-0.5B, ViT-H/14|
-| [`uform-gen`][model-g1]             |    1.5B    | Image Captioning, VQA       | llama-1.3B, ViT-B/16  |
+| Model                              | Parameters |                     Purpose |           Architecture |
+| :--------------------------------- | ---------: | --------------------------: | ---------------------: |
+| [`uform-gen2-dpo`][model-g2] 🆕     |       1.2B | Chat, Image Captioning, VQA | qwen1.5-0.5B, ViT-H/14 |
+| [`uform-gen2-qwen-500m`][model-g2] |       1.2B | Chat, Image Captioning, VQA | qwen1.5-0.5B, ViT-H/14 |
+| [`uform-gen`][model-g1]            |       1.5B |       Image Captioning, VQA |   llama-1.3B, ViT-B/16 |
 
 [model-g2]: https://huggingface.co/unum-cloud/uform-gen2-qwen-500m/
 [model-g1]: https://huggingface.co/unum-cloud/uform-gen/
 
+## Producing Embeddings
 
-## Quick Start
-
-Once you `pip install uform`, fetching the models is as easy as:
+Add UForm to your dependencies list, or just install it locally:
 
-```python
-import uform
-
-model = uform.get_model('unum-cloud/uform-vl-english') # Just English
-model = uform.get_model('unum-cloud/uform-vl-multilingual-v2') # 21 Languages
+```bash
+pip install uform
 ```
 
-### Producing Embeddings
+Then, you can use the following code to get embeddings for text and images.
+You can do that either with the PyTorch reference model or the lighter cross-platform ONNX weights.
 
 ```python
+import uform
 from PIL import Image
-import torch.nn.functional as F
+
+# If you want to use the PyTorch model
+model, processor = uform.get_model('unum-cloud/uform-vl-english-large') # Just English
+model, processor = uform.get_model('unum-cloud/uform-vl-multilingual-v2') # 21 Languages
+
+# If you want to use the light-weight portable ONNX model
+# Available combinations: cpu & fp32, gpu & fp32, gpu & fp16
+# Check out Unum's Hugging Face space for more details: https://huggingface.co/unum-cloud
+model, processor = uform.get_model_onnx('unum-cloud/uform-vl-english-small', 'cpu', 'fp32')
+model, processor = uform.get_model_onnx('unum-cloud/uform-vl-english-large', 'gpu', 'fp16')
 
 text = 'a small red panda in a zoo'
 image = Image.open('red_panda.jpg')
 
-image_data = model.preprocess_image(image)
-text_data = model.preprocess_text(text)
+image_data = processor.preprocess_image(image)
+text_data = processor.preprocess_text(text)
 
 image_features, image_embedding = model.encode_image(image_data, return_features=True)
 text_features, text_embedding = model.encode_text(text_data, return_features=True)
-
-similarity = F.cosine_similarity(image_embedding, text_embedding)
 ```
 
 To search for similar items, the embeddings can be compared using cosine similarity.
 The resulting value will fall within the range of `-1` to `1`, where `1` indicates a high likelihood of a match.
+PyTorch provides a built-in function for calculating cosine similarity, while for ONNX, you can use NumPy.
+
+```python
+import torch.nn.functional as F
+
+similarity = F.cosine_similarity(image_embedding, text_embedding)
+```
+
+ONNX has no such function, but you can calculate the cosine similarity using [SimSIMD](https://github.com/ashvardanian/simsimd) or manually, with NumPy:
+
+```python
+import numpy as np
+
+image_embedding = image_embedding / np.linalg.norm(image_embedding, keepdims=True, axis=1)
+text_embedding = text_embedding / np.linalg.norm(text_embedding, keepdims=True, axis=1)
+similarity = (image_embedding * text_embedding).sum(axis=1)
+```
+
+### Reranking
+
 Once the list of nearest neighbors (best matches) is obtained, the joint multimodal embeddings, created from both text and image features, can be used to better rerank (reorder) the list.
 The model can calculate a "matching score" that falls within the range of `[0, 1]`, where `1` indicates a high likelihood of a match.
 
 ```python
-joint_embedding = model.encode_multimodal(
+score, joint_embedding = model.encode_multimodal(
     image_features=image_features,
     text_features=text_features,
-    attention_mask=text_data['attention_mask']
+    attention_mask=text_data['attention_mask'],
+    return_scores=True,
 )
-score = model.get_matching_scores(joint_embedding)
 ```
 
-### Chat, Image Captioning and Question Answering
+### Down-casting, Quantization, Matryoshka, and Slicing
+
+Depending on the application, the embeddings can be down-casted to smaller numeric representations without losing much recall.
+Switching from `f32` to `f16` is recommended in almost all cases, unless you are running on very old hardware without half-precision support.
+Switching to `i8` with linear scaling is also possible, but will be noticeable in the recall on larger collections with millions of searchable entries.
+Similarly, for higher-dimensional embeddings (512 or 768), a common strategy is to quantize them into single-bit representations for faster search.
+
+```python
+import numpy as np
+
+f32_embedding: np.ndarray = model.encode_text(text_data, return_features=False).detach().cpu().numpy()
+f16_embedding: np.ndarray = f32_embedding.astype(np.float16)
+i8_embedding: np.ndarray = (f32_embedding * 127).astype(np.int8)
+b1_embedding: np.ndarray = np.packbits((f32_embedding > 0).astype(np.uint8))
+```
+
+Alternative approach to quantization is to use the Matryoshka embeddings, where the embeddings are sliced into smaller parts, and the search is performed in a hierarchical manner.
+
+```python
+import numpy as np
+
+large_embedding: np.ndarray = model.encode_text(text_data, return_features=False).detach().cpu().numpy()
+small_embedding: np.ndarray = large_embedding[:, :256]
+tiny_embedding: np.ndarray = large_embedding[:, :64]
+```
+
+Both approaches are natively supported by the [USearch][github-usearch] vector-search engine and the [SimSIMD][github-simsimd] numerics libraries.
+When dealing with small collections (up to millions of entries) and looking for low-latency cosine distance calculations, you can [achieve 5x-2500x performance improvement][report-simsimd] over Torch, NumPy, SciPy, and vanilla Python using SimSIMD.
+
+```python
+from simsimd import cosine, hamming
+
+distance: float = cosine(f32_embedding, f32_embedding) # 32x SciPy performance on Apple M2 CPU
+distance: float = cosine(f16_embedding, f16_embedding) # 79x SciPy performance on Apple M2 CPU
+distance: float = cosine(i8_embedding, i8_embedding) # 133x SciPy performance on Apple M2 CPU
+distance: float = hamming(b1_embedding, b1_embedding) # 17x SciPy performance on Apple M2 CPU
+```
+
+Similarly, when dealing with large collections (up to billions of entries per server) and looking for high-throughput search, you can [achieve 100x performance improvement][report-usearch] over FAISS and other vector-search solutions using USearch.
+Here are a couple of examples:
+
+```python
+from usearch.index import Index
+
+f32_index = Index(ndim=64, metric='cos', dtype='f32') # for Matryoshka embeddings
+f16_index = Index(ndim=64, metric='cos', dtype='f16') # for Matryoshka embeddings
+i8_index = Index(ndim=256, metric='cos', dtype='i8') # for quantized embeddings
+b1_index = Index(ndim=768, metric='hamming', dtype='b1') # for binary embeddings
+```
+
+[github-usearch]: https://github.com/unum-cloud/usearch
+[github-simsimd]: https://github.com/ashvardanian/simsimd
+[report-usearch]: https://www.unum.cloud/blog/2023-11-07-scaling-vector-search-with-intel
+[report-simsimd]: https://ashvardanian.com/posts/python-c-assembly-comparison/
+
+### Compact Packaging
+
+PyTorch is a heavy dependency to carry, especially if you run on Edge or IoT devices.
+Using vanilla ONNX runtime, one can significantly reduce memory consumption and deployment latency.
+
+```sh
+$ conda create -n uform_torch python=3.10 -y
+$ conda create -n uform_onnx python=3.10 -y
+$ conda activate uform_torch && pip install -e ".[torch]" && conda deactivate
+$ conda activate uform_onnx && pip install -e ".[onnx]" && conda deactivate
+$ du -sh $(conda info --envs | grep 'uform_torch' | awk '{print $2}')
+> 5.2G    ~/conda/envs/uform_torch
+$ du -sh $(conda info --envs | grep 'uform_onnx' | awk '{print $2}')
+> 461M    ~/conda/envs/uform_onnx
+```
+
+Most of that weight can be further reduced down to 100 MB for both the model and the runtime.
+You can pick one of many supported [ONNX execution providers][onnx-providers], which includes XNNPACK, CUDA and TensorRT for Nvidia GPUs, OpenVINO on Intel, DirectML on Windows, ROCm on AMD, CoreML on Apple devices, and more to come.
+
+[onnx-providers]: https://onnxruntime.ai/docs/execution-providers/
+
+---
+
+The configuration process may include a few additional steps, depending on the environment.
+When using the CUDA and TensorRT backends with CUDA 12 or newer make sure to [install the Nvidia toolkit][install-nvidia-toolkit] and the `onnxruntime-gpu` package from the custom repository.
+
+```sh
+wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
+sudo dpkg -i cuda-keyring_1.1-1_all.deb
+sudo apt-get update
+sudo apt-get -y install cuda-toolkit-12
+pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
+export CUDA_PATH="/usr/local/cuda-12/bin"
+export PATH="/usr/local/cuda-12/bin${PATH:+:${PATH}}"
+export LD_LIBRARY_PATH="/usr/local/cuda-12/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
+pytest python/scripts/ -s -x -Wd -v -k onnx
+```
+
+[install-nvidia-toolkit]: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/#network-repo-installation-for-ubuntu
 
-The generative model can be used to caption images, answer questions about them. Also it is suitable for a multimodal chat.
+## Chat, Image Captioning and Question Answering
 
+UForm generative models are fully compatible with the Hugging Face Transformers library, and can be used without installing the UForm library.
+Those models can be used to caption images or power multimodal chat experiences.
 
 ```python
 from transformers import AutoModel, AutoProcessor
 
-model = AutoModel.from_pretrained("unum-cloud/uform-gen2-qwen-500m", trust_remote_code=True)
-processor = AutoProcessor.from_pretrained("unum-cloud/uform-gen2-qwen-500m", trust_remote_code=True)
+model = AutoModel.from_pretrained('unum-cloud/uform-gen2-qwen-500m', trust_remote_code=True)
+processor = AutoProcessor.from_pretrained('unum-cloud/uform-gen2-qwen-500m', trust_remote_code=True)
 
-prompt = "Question or Instruction"
-image = Image.open("image.jpg")
+prompt = 'Question or Instruction'
+image = Image.open('image.jpg')
 
-inputs = processor(text=[prompt], images=[image], return_tensors="pt")
+inputs = processor(text=[prompt], images=[image], return_tensors='pt')
 
 with torch.inference_mode():
      output = model.generate(
@@ -130,7 +263,7 @@ with torch.inference_mode():
         eos_token_id=151645,
         pad_token_id=processor.tokenizer.pad_token_id
     )
-prompt_len = inputs["input_ids"].shape[1]
+prompt_len = inputs['input_ids'].shape[1]
 decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
 ```
 
@@ -148,16 +281,16 @@ The exact behavior is controlled by prompts.
 ```python
 from uform.gen_model import VLMForCausalLM, VLMProcessor
 
-model = VLMForCausalLM.from_pretrained("unum-cloud/uform-gen")
-processor = VLMProcessor.from_pretrained("unum-cloud/uform-gen")
+model = VLMForCausalLM.from_pretrained('unum-cloud/uform-gen')
+processor = VLMProcessor.from_pretrained('unum-cloud/uform-gen')
 
 # [cap] Narrate the contents of the image with precision.
 # [cap] Summarize the visual content of the image.
 # [vqa] What is the main subject of the image?
-prompt = "[cap] Summarize the visual content of the image."
-image = Image.open("zebra.jpg")
+prompt = '[cap] Summarize the visual content of the image.'
+image = Image.open('zebra.jpg')
 
-inputs = processor(texts=[prompt], images=[image], return_tensors="pt")
+inputs = processor(texts=[prompt], images=[image], return_tensors='pt')
 with torch.inference_mode():
      output = model.generate(
         **inputs,
@@ -168,7 +301,7 @@ with torch.inference_mode():
         pad_token_id=processor.tokenizer.pad_token_id
     )
 
-prompt_len = inputs["input_ids"].shape[1]
+prompt_len = inputs['input_ids'].shape[1]
 decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
 ```
 
@@ -193,10 +326,10 @@ For that pick the encoder of the model you want to run in parallel (`text_encode
 ```python
 import uform
 
-model = uform.get_model('unum-cloud/uform-vl-english')
+model, processor = uform.get_model('unum-cloud/uform-vl-english')
 model_image = nn.DataParallel(model.image_encoder)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model_image.to(device)
 
 _, res = model_image(images, 0)
@@ -266,11 +399,11 @@ Evaluating the `unum-cloud/uform-vl-multilingual-v2` model, one can expect the f
 
 ### Generative Models
 
-| Model                               | LLM Size |  SQA  |  MME   | MMBench  | Average¹ |
-| :---------------------------------- | -------: | -----:| ------:| --------:| --------:|
-| UForm-Gen2-Qwen-500m                |   0.5B   | 45.5  | 880.1  |  42.0    |   29.31  |
-| MobileVLM v2                        |   1.4B   | 52.1  | 1302.8 |  57.7    |   36.81  |
-| LLaVA-Phi                           |   2.7B   | 68.4  | 1335.1 |  59.8    |   42.95  |
+| Model                | LLM Size |  SQA |    MME | MMBench | Average¹ |
+| :------------------- | -------: | ---: | -----: | ------: | -------: |
+| UForm-Gen2-Qwen-500m |     0.5B | 45.5 |  880.1 |    42.0 |    29.31 |
+| MobileVLM v2         |     1.4B | 52.1 | 1302.8 |    57.7 |    36.81 |
+| LLaVA-Phi            |     2.7B | 68.4 | 1335.1 |    59.8 |    42.95 |
 
 For captioning evaluation we measure CLIPScore and RefCLIPScore³.
 
diff --git a/VERSION b/VERSION
index d15723f..45a1b3f 100644
--- a/VERSION
+++ b/VERSION
@@ -1,1 +1,1 @@
-0.3.2
+1.1.2
diff --git a/package-ci.json b/package-ci.json
new file mode 100644
index 0000000..e5e5b93
--- /dev/null
+++ b/package-ci.json
@@ -0,0 +1,10 @@
+{
+    "name": "uform-ci",
+    "version": "1.0.0",
+    "devDependencies": {
+        "@semantic-release/exec": "^6.0.3",
+        "@semantic-release/git": "^10.0.1",
+        "conventional-changelog-eslint": "^3.0.9",
+        "semantic-release": "^21.1.2"
+    }
+}
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index dc322da..0837c51 100755
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,25 +31,33 @@ classifiers = [
 dependencies = [
     "huggingface_hub>=0.16.4",
     "tokenizers>=0.13.3",
-    "torch>=1.13.1",
-    "torchvision",
-    "transformers>=4.36.2",
+    "pillow"
 ]
-description = "Multi-Modal Transformers library for Semantic Search and other Vision-Language tasks"
+description = "Pocket-Sized Multimodal AI for Content Understanding and Generation"
 maintainers = [
     {email = "info@unum.cloud", name = "Unum Cloud"},
 ]
 name = "uform"
 readme = "README.md"
 requires-python = ">=3.7"
-version = "1.1.1"
+version = "1.1.2"
 
 [project.scripts]
 uform-chat = "uform.chat:main"
 
+[project.optional-dependencies]
+torch = ["torch>=1.13.1", "torchvision", "transformers>=4.36.2"]
+onnx = ["onnx>=1.15.0", "onnxruntime>=1.17.1", "numpy"]
+onnx-gpu = ["onnx>=1.15.0", "onnxruntime-gpu>=1.17.1", "numpy"]
+
 [project.urls]
 "Homepage" = "https://github.com/unum-cloud/uform"
 
+[tool.setuptools.packages.find]
+where = ["python"]
+include = ["uform"]
+namespaces = false
+
 [tool.ruff]
 ignore = ["C408", "C901", "E501", "E741"]
 ignore-init-module-imports = true
@@ -70,3 +78,9 @@ in_place = true
 spaces_before_inline_comment = 2
 spaces_indent_inline_array = 4
 trailing_comma_inline_array = true
+
+# Configuration options for the Black formatter:
+# https://black.readthedocs.io/en/latest/usage_and_configuration/the_basics.html#where-black-looks-for-the-file
+[tool.black]
+line-length = 120                   # Set line length to the same value as in `.clang-format` for modern wide screens
+target-version = ['py36', 'py312']  # Set target Python versions to 3.6 and 3.12
\ No newline at end of file
diff --git a/scripts/bench.py b/python/scripts/bench.py
similarity index 96%
rename from scripts/bench.py
rename to python/scripts/bench.py
index ed61589..49c7004 100644
--- a/scripts/bench.py
+++ b/python/scripts/bench.py
@@ -80,9 +80,7 @@ def bench_image_embeddings(model, images):
     total_embeddings = 0
     images *= 10
     while total_duration < 10:
-        seconds, embeddings = duration(
-            lambda: model.encode_image(model.preprocess_image(images))
-        )
+        seconds, embeddings = duration(lambda: model.encode_image(processor.preprocess_image(images)))
         total_duration += seconds
         total_embeddings += len(embeddings)
 
@@ -94,9 +92,7 @@ def bench_text_embeddings(model, texts):
     total_embeddings = 0
     texts *= 10
     while total_duration < 10:
-        seconds, embeddings = duration(
-            lambda: model.encode_text(model.preprocess_text(texts))
-        )
+        seconds, embeddings = duration(lambda: model.encode_text(processor.preprocess_text(texts)))
         total_duration += seconds
         total_embeddings += len(embeddings)
 
diff --git a/python/uform/__init__.py b/python/uform/__init__.py
new file mode 100755
index 0000000..1ecb242
--- /dev/null
+++ b/python/uform/__init__.py
@@ -0,0 +1,59 @@
+from json import load
+from os.path import join
+from typing import Mapping, Optional, Tuple
+
+from huggingface_hub import snapshot_download
+
+
+def get_checkpoint(model_name: str, token: str) -> Tuple[str, Mapping, str]:
+    import torch
+
+    model_path = snapshot_download(repo_id=model_name, token=token)
+    config_path = join(model_path, "torch_config.json")
+
+    state = torch.load(join(model_path, "torch_weight.pt"))
+    return config_path, state, join(model_path, "tokenizer.json")
+
+
+def get_model(model_name: str, token: Optional[str] = None):
+    from uform.torch_models import VLM
+    from uform.torch_preprocessor import TorchProcessor
+
+    config_path, state, tokenizer_path = get_checkpoint(model_name, token)
+
+    with open(config_path) as f:
+        config = load(f)
+
+    model = VLM(config, tokenizer_path)
+    model.image_encoder.load_state_dict(state["image_encoder"])
+    model.text_encoder.load_state_dict(state["text_encoder"])
+    processor = TorchProcessor(config, tokenizer_path)
+
+    return model.eval(), processor
+
+
+def get_model_onnx(model_name: str, device: str, dtype: str, token: Optional[str] = None):
+    from uform.onnx_models import VLM_ONNX
+    from uform.numpy_preprocessor import NumPyProcessor
+
+    assert device in (
+        "cpu",
+        "gpu",
+    ), f"Invalid `device`: {device}. Must be either `cpu` or `gpu`"
+    assert dtype in (
+        "fp32",
+        "fp16",
+    ), f"Invalid `dtype`: {dtype}. Must be either `fp32` or `fp16` (only for gpu)"
+    assert (
+        device == "cpu" and dtype == "fp32"
+    ) or device == "gpu", "Combination `device`=`cpu` & `dtype=fp16` is not supported"
+
+    model_path = snapshot_download(repo_id=f"{model_name}-{device}-{dtype}", token=token)
+
+    with open(join(model_path, "config.json")) as f:
+        config = load(f)
+
+    model = VLM_ONNX(model_path, config, device, dtype)
+    processor = NumPyProcessor(config, join(model_path, "tokenizer.json"))
+
+    return model, processor
diff --git a/src/uform/chat.py b/python/uform/chat.py
similarity index 100%
rename from src/uform/chat.py
rename to python/uform/chat.py
diff --git a/src/uform/gen_model.py b/python/uform/gen_model.py
similarity index 97%
rename from src/uform/gen_model.py
rename to python/uform/gen_model.py
index 63905ce..35faae1 100644
--- a/src/uform/gen_model.py
+++ b/python/uform/gen_model.py
@@ -3,20 +3,15 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from torchvision.transforms import (
-    CenterCrop,
-    Compose,
-    InterpolationMode,
-    Normalize,
-    RandomResizedCrop,
-    Resize,
-    ToTensor,
-)
+from torchvision.transforms import (CenterCrop, Compose, InterpolationMode,
+                                    Normalize, RandomResizedCrop, Resize,
+                                    ToTensor)
 from transformers import AutoConfig, AutoTokenizer
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
-from transformers.models.auto.modeling_auto import AutoModel, AutoModelForCausalLM
+from transformers.models.auto.modeling_auto import (AutoModel,
+                                                    AutoModelForCausalLM)
 from transformers.processing_utils import ProcessorMixin
 from transformers.tokenization_utils_base import BatchEncoding
 
diff --git a/python/uform/numpy_preprocessor.py b/python/uform/numpy_preprocessor.py
new file mode 100644
index 0000000..6cdd54b
--- /dev/null
+++ b/python/uform/numpy_preprocessor.py
@@ -0,0 +1,94 @@
+from os import PathLike
+from typing import Dict, List, Union
+
+from PIL.Image import Image, BICUBIC
+from tokenizers import Tokenizer
+import numpy as np
+
+
+class NumPyProcessor:
+    def __init__(self, config: Dict, tokenizer_path: PathLike):
+        """
+        :param config: model config
+        :param tokenizer_path: path to tokenizer file
+        :param tensor_type: which tensors to return, either pt (PyTorch) or np (NumPy)
+        """
+
+        self._image_size = config["image_encoder"]["image_size"]
+        self._max_seq_len = config["text_encoder"]["max_position_embeddings"]
+        self._tokenizer = Tokenizer.from_file(tokenizer_path)
+        self._tokenizer.no_padding()
+        self._pad_token_idx = config["text_encoder"]["padding_idx"]
+
+        self.image_mean = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)[None, None]
+        self.image_std = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)[None, None]
+
+    def preprocess_text(self, texts: Union[str, List[str]]) -> Dict[str, np.ndarray]:
+        """Transforms one or more strings into dictionary with tokenized strings and attention masks.
+
+        :param texts: text of list of texts to tokenizer
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+
+        input_ids = np.full(
+            (len(texts), self._max_seq_len),
+            fill_value=self._pad_token_idx,
+            dtype=np.int64,
+        )
+
+        attention_mask = np.zeros(
+            (len(texts), self._max_seq_len),
+            dtype=np.int32,
+        )
+        encoded = self._tokenizer.encode_batch(texts)
+
+        for i, seq in enumerate(encoded):
+            seq_len = min(len(seq), self._max_seq_len)
+            input_ids[i, :seq_len] = seq.ids[:seq_len]
+
+            attention_mask[i, :seq_len] = 1
+
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+    def preprocess_image(self, images: Union[Image, List[Image]]) -> np.ndarray:
+        """Transforms one or more Pillow images into Torch Tensors.
+
+        :param images: image or list of images to preprocess
+        """
+
+        if isinstance(images, list):
+            batch_images = np.empty(
+                (len(images), 3, self._image_size, self._image_size),
+                dtype=np.float32,
+            )
+
+            for i, image in enumerate(images):
+                batch_images[i] = self._resize_crop_normalize(image)
+
+        else:
+            batch_images = self._resize_crop_normalize(images)[None]
+
+        return batch_images
+
+    def _resize_crop_normalize(self, image: Image):
+        width, height = image.size
+
+        if width < height:
+            width = self._image_size
+            height = int(height / width * self._image_size)
+        else:
+            width = int(width / height * self._image_size)
+            height = self._image_size
+
+        image = image.resize((width, height), resample=BICUBIC)
+
+        left = (width - self._image_size) / 2
+        top = (height - self._image_size) / 2
+        right = (width + self._image_size) / 2
+        bottom = (height + self._image_size) / 2
+
+        image = image.convert("RGB").crop((left, top, right, bottom))
+        image = (np.array(image).astype(np.float32) / 255.0 - self.image_mean) / self.image_std
+
+        return np.transpose(image, (2, 0, 1))
diff --git a/python/uform/onnx_models.py b/python/uform/onnx_models.py
new file mode 100644
index 0000000..8e2a87a
--- /dev/null
+++ b/python/uform/onnx_models.py
@@ -0,0 +1,231 @@
+from os.path import join
+from typing import Dict, Optional, Tuple, Union
+
+import onnxruntime as ort
+from numpy import ndarray
+
+
+class ExecutionProviderError(Exception):
+    """Exception raised when a requested execution provider is not available."""
+
+
+def available_providers(device: str) -> Tuple[str, ...]:
+    gpu_providers = ("CUDAExecutionProvider", "TensorrtExecutionProvider")
+    cpu_providers = ("OpenVINOExecutionProvider", "CoreMLExecutionProvider", "CPUExecutionProvider")
+    available = ort.get_available_providers()
+    if device == "gpu":
+        if all(provider not in available for provider in gpu_providers):
+            raise ExecutionProviderError(
+                f"GPU providers are not available, consider installing `onnxruntime-gpu` and make sure the CUDA is available on your system. Currently installed: {available}"
+            )
+        return gpu_providers
+
+    return cpu_providers
+
+
+class VisualEncoderONNX:
+    def __init__(self, model_path: str, device: str):
+        """
+        :param model_path: Path to onnx model
+        :param device: Device name, either cpu or gpu
+        """
+
+        session_options = ort.SessionOptions()
+        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        self.session = ort.InferenceSession(
+            model_path,
+            sess_options=session_options,
+            providers=available_providers(device),
+        )
+
+    def __call__(self, images: ndarray) -> Tuple[ndarray, ndarray]:
+        return self.session.run(None, {"images": images})
+
+
+class TextEncoderONNX:
+    def __init__(self, text_encoder_path: str, reranker_path: str, device: str):
+        """
+        :param text_encoder_path: Path to onnx of text encoder
+        :param reranker_path: Path to onnx of reranker
+        :param device: Device name, either cpu or gpu
+        """
+
+        session_options = ort.SessionOptions()
+        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        self.text_encoder_session = ort.InferenceSession(
+            text_encoder_path,
+            sess_options=session_options,
+            providers=available_providers(device),
+        )
+
+        self.reranker_session = ort.InferenceSession(
+            reranker_path,
+            sess_options=session_options,
+            providers=available_providers(device),
+        )
+
+    def __call__(self, input_ids: ndarray, attention_mask: ndarray) -> Tuple[ndarray, ndarray]:
+        return self.text_encoder_session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
+
+    def forward_multimodal(
+        self, text_features: ndarray, attention_mask: ndarray, image_features: ndarray
+    ) -> Tuple[ndarray, ndarray]:
+        return self.reranker_session.run(
+            None,
+            {
+                "text_features": text_features,
+                "attention_mask": attention_mask,
+                "image_features": image_features,
+            },
+        )
+
+
+class VLM_ONNX:
+    def __init__(self, checkpoint_path: str, config: Dict, device: str, dtype: str):
+        assert device in (
+            "cpu",
+            "gpu",
+        ), f"Invalid `device`: {device}. Must be either `cpu` or `gpu`"
+        assert dtype in (
+            "fp32",
+            "fp16",
+        ), f"Invalid `dtype`: {dtype}. Must be either `fp32` or `fp16` (only for gpu)"
+        assert (
+            device == "cpu" and dtype == "fp32"
+        ) or device == "gpu", "Combination `device`=`cpu` & `dtype=fp16` is not supported"
+
+        self.device = device
+        self.dtype = dtype
+
+        self._embedding_dim = config["text_encoder"]["embedding_dim"]
+        self._text_encoder_dim = config["text_encoder"]["dim"]
+        self._image_encoder_dim = config["image_encoder"]["dim"]
+
+        self.text_encoder = TextEncoderONNX(
+            join(checkpoint_path, f"text_encoder.onnx"),
+            join(checkpoint_path, f"reranker.onnx"),
+            device,
+        )
+
+        self.image_encoder = VisualEncoderONNX(join(checkpoint_path, f"image_encoder.onnx"), device)
+
+    def encode_image(
+        self,
+        images: ndarray,
+        return_features: bool = False,
+    ) -> Union[ndarray, Tuple[ndarray, ndarray]]:
+        """Passes the pre-processed images through `image_encoder` to produce images features (optional) and embeddings.
+
+        :param images: Preprocessed image
+        :param return_features: Whether to return images features or return only embeddings
+        """
+
+        features, embeddings = self.image_encoder(images)
+
+        if return_features:
+            return features, embeddings
+
+        return embeddings
+
+    def encode_text(
+        self,
+        texts: Dict[str, ndarray],
+        return_features: bool = False,
+    ) -> Union[ndarray, Tuple[ndarray, ndarray]]:
+        """Passes the pre-processed texts through `text_encoder` to produce texts features (optional) and embeddings.
+
+        :param texts: Dictionary with tokenized texts and attention masks
+        :param return_features: Whether to return texts features or return only embeddings
+        """
+
+        features, embeddings = self.text_encoder(**texts)
+
+        if return_features:
+            return features, embeddings
+
+        return embeddings
+
+    def encode_multimodal(
+        self,
+        image: Optional[ndarray] = None,
+        text: Dict[str, ndarray] = None,
+        image_features: Optional[ndarray] = None,
+        text_features: Optional[ndarray] = None,
+        attention_mask: Optional[ndarray] = None,
+        return_scores: bool = False,
+    ) -> Union[ndarray, Tuple[ndarray, ndarray]]:
+        """Passes preprocessed texts (or precomputed texts features) and
+            preprocessed images (or precomputed images features) through multimodal encoded to produce matching scores and optionally multimodal joint embeddings.
+
+        :param image: Preprocessed images
+        :param text: Preprocessed texts
+        :param image_features: Precomputed images features
+        :param text_features: Precomputed text features
+        :param attention_mask: Attention masks, not required if pass `text` instead of text_features
+        """
+
+        assert image is not None or image_features is not None, "Either `image` or `image_features` should be non None"
+        assert text is not None or text_features is not None, "Either `text_data` or `text_features` should be non None"
+
+        if text_features is not None:
+            assert attention_mask is not None, "if `text_features` is not None, then you should pass `attention_mask`"
+
+        if image_features is None:
+            image_features = self.image_encoder(image)
+
+        if text_features is None:
+            text_features = self.text_encoder(
+                text["input_ids"],
+                text["attention_mask"],
+            )
+
+        matching_scores, embeddings = self.text_encoder.forward_multimodal(
+            text_features,
+            attention_mask if attention_mask is not None else text["attention_mask"],
+            image_features,
+        )
+
+        if return_scores:
+            return matching_scores, embeddings
+
+        return embeddings
+
+    def forward(
+        self,
+        images: ndarray,
+        texts: Dict[str, ndarray],
+    ) -> Union[ndarray, ndarray]:
+        """Inference forward method
+
+        :param images: Preprocessed images
+        :param texts: Preprocessed texts
+        :return: embeddings for images and texts
+        """
+        _, image_embeddings = self.image_encoder(images)
+        _, text_embeddings = self.text_encoder(texts)
+        return image_embeddings, text_embeddings
+
+    @property
+    def text_features_dim(self) -> int:
+        """Dimensionality of the text encoder features."""
+
+        return self._text_encoder_dim
+
+    @property
+    def image_features_dim(self) -> int:
+        """Dimensionality of the image encoder features."""
+
+        return self._image_encoder_dim
+
+    @property
+    def embedding_dim(self) -> int:
+        """Dimensionality of shared space embedding."""
+
+        return self._embedding_dim
+
+    @property
+    def multimodal_embedding_dim(self) -> int:
+        """Dimensionality of multimodal joint embedding."""
+        return self._text_encoder_dim
diff --git a/python/uform/preprocessing.py b/python/uform/preprocessing.py
new file mode 100644
index 0000000..d3d833e
--- /dev/null
+++ b/python/uform/preprocessing.py
@@ -0,0 +1,105 @@
+from os import PathLike
+from typing import Dict, List, Union
+
+import torch
+from PIL import Image
+from tokenizers import Tokenizer
+from torch import Tensor
+from torchvision.transforms import (CenterCrop, Compose, InterpolationMode,
+                                    Normalize, Resize, ToTensor)
+
+
+# lambda is not pickable
+def convert_to_rgb(image):
+    return image.convert("RGB")
+
+
+class Processor:
+    def __init__(self, config: Dict, tokenizer_path: PathLike, tensor_type: str = "pt"):
+        """
+        :param config: model config
+        :param tokenizer_path: path to tokenizer file
+        :param tensor_type: which tensors to return, either pt (PyTorch) or np (NumPy)
+        """
+
+        assert tensor_type in ("pt", "np"), "`tensor_type` must be either `pt` or `np`"
+
+        self._image_size = config["image_encoder"]["image_size"]
+        self._max_seq_len = config["text_encoder"]["max_position_embeddings"]
+        self._tokenizer = Tokenizer.from_file(tokenizer_path)
+        self._tokenizer.no_padding()
+        self._pad_token_idx = config["text_encoder"]["padding_idx"]
+
+        self.tensor_type = tensor_type
+
+        self._image_transform = Compose(
+            [
+                Resize(self._image_size, interpolation=InterpolationMode.BICUBIC),
+                convert_to_rgb,
+                CenterCrop(self._image_size),
+                ToTensor(),
+                Normalize(
+                    mean=(0.48145466, 0.4578275, 0.40821073),
+                    std=(0.26862954, 0.26130258, 0.27577711),
+                ),
+            ],
+        )
+
+    def preprocess_text(self, texts: Union[str, List[str]]) -> Dict[str, Tensor]:
+        """Transforms one or more strings into dictionary with tokenized strings and attention masks.
+
+        :param texts: text of list of texts to tokenizer
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+
+        input_ids = torch.full(
+            (len(texts), self._max_seq_len),
+            fill_value=self._pad_token_idx,
+            dtype=torch.int64,
+        )
+
+        attention_mask = torch.zeros(
+            len(texts),
+            self._max_seq_len,
+            dtype=torch.int32,
+        )
+        encoded = self._tokenizer.encode_batch(texts)
+
+        for i, seq in enumerate(encoded):
+            seq_len = min(len(seq), self._max_seq_len)
+            input_ids[i, :seq_len] = torch.LongTensor(
+                seq.ids[: self._max_seq_len],
+            )
+            attention_mask[i, :seq_len] = 1
+
+        if self.tensor_type == "np":
+            return {
+                "input_ids": input_ids.numpy(),
+                "attention_mask": attention_mask.numpy(),
+            }
+
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+    def preprocess_image(self, images: Union[Image, List[Image]]) -> Tensor:
+        """Transforms one or more Pillow images into Torch Tensors.
+
+        :param images: image or list of images to preprocess
+        """
+
+        if isinstance(images, list):
+            batch_images = torch.empty(
+                (len(images), 3, self._image_size, self._image_size),
+                dtype=torch.float32,
+            )
+
+            for i, image in enumerate(images):
+                batch_images[i] = self._image_transform(image)
+
+        else:
+            batch_images = self._image_transform(images).unsqueeze(0)
+
+        if self.tensor_type == "np":
+            return batch_images.numpy()
+
+        return batch_images
diff --git a/src/uform/models.py b/python/uform/torch_models.py
similarity index 58%
rename from src/uform/models.py
rename to python/uform/torch_models.py
index d9c8fa7..d0d810b 100644
--- a/src/uform/models.py
+++ b/python/uform/torch_models.py
@@ -1,26 +1,11 @@
 from dataclasses import dataclass
 from os import PathLike
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from PIL.Image import Image
-from tokenizers import Tokenizer
 from torch import Tensor
-from torchvision.transforms import (
-    CenterCrop,
-    Compose,
-    InterpolationMode,
-    Normalize,
-    Resize,
-    ToTensor,
-)
-
-
-# lambda is not pickable
-def convert_to_rgb(image):
-    return image.convert("RGB")
 
 
 @dataclass(eq=False)
@@ -273,9 +258,7 @@ def get_attention_mask(self, attn_mask: Tensor, dtype: torch.dtype) -> Tensor:
     def get_position_ids(self, x: Tensor) -> Tensor:
         if self.model_type == "roberta":
             mask = x.ne(self.padding_idx).int()
-            return (
-                torch.cumsum(mask, dim=1).type_as(mask) * mask
-            ).long() + self.padding_idx
+            return (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + self.padding_idx
 
         return self.position_ids[:, : x.shape[1]]
 
@@ -284,7 +267,7 @@ def embed_text(self, x: Tensor) -> Tensor:
         x = self.word_embeddings(x) + positional_embedding
         return self.dropout(self.layer_norm(x))
 
-    def forward(self, x: dict) -> torch.Tensor:
+    def forward(self, x: dict) -> Tensor:
         features = self.forward_features(x["input_ids"], x["attention_mask"])
         embeddings = self.forward_embedding(features, x["attention_mask"])
         return features, embeddings
@@ -299,6 +282,7 @@ class VisualEncoder(nn.Module):
     num_heads: int
     embedding_dim: int
     pooling: str
+    num_reg_tokens: int = 0
 
     def __post_init__(self):
         super().__init__()
@@ -308,11 +292,11 @@ def __post_init__(self):
         self.pos_embed = nn.Parameter(torch.randn(1, seq_len, self.dim) * 0.02)
         self.cls_token = nn.Parameter(torch.zeros(1, 1, self.dim))
 
+        if self.num_reg_tokens > 0:
+            self.reg_token = nn.Parameter(torch.zeros(1, self.num_reg_tokens, self.dim))
+
         self.blocks = nn.Sequential(
-            *[
-                VisualEncoderBlock(self.dim, self.num_heads)
-                for _ in range(self.num_layers)
-            ],
+            *[VisualEncoderBlock(self.dim, self.num_heads) for _ in range(self.num_layers)],
         )
 
         self.norm = nn.LayerNorm(self.dim, eps=1e-6)
@@ -321,7 +305,14 @@ def __post_init__(self):
     def forward_features(self, x: Tensor) -> Tensor:
         x = self.patch_embed(x).flatten(start_dim=2).transpose(2, 1)
         x = x + self.pos_embed
-        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+
+        special_tokens = [self.cls_token.expand(x.shape[0], -1, -1)]
+
+        if self.num_reg_tokens > 0:
+            special_tokens.append(self.reg_token.expand(x.shape[0], -1, -1))
+
+        x = torch.cat(special_tokens + [x], dim=1)
+
         x = self.blocks(x)
 
         return self.norm(x)
@@ -334,7 +325,7 @@ def forward_embedding(self, x: Tensor) -> Tensor:
 
         return self.embedding_projection(x)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         features = self.forward_features(x)
         embeddings = self.forward_embedding(features)
         return features, embeddings
@@ -342,7 +333,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class VLM(nn.Module):
     """
-    Vision-Language Model for multi-modal embeddings.
+    Vision-Language Model for Multimodal embeddings.
     """
 
     def __init__(self, config: Dict, tokenizer_path: PathLike):
@@ -351,30 +342,11 @@ def __init__(self, config: Dict, tokenizer_path: PathLike):
         """
 
         super().__init__()
-        self._max_seq_len = config["text_encoder"]["max_position_embeddings"]
         self._embedding_dim = config["text_encoder"]["embedding_dim"]
-        self._image_size = config["image_encoder"]["image_size"]
 
         self.text_encoder = TextEncoder(**config["text_encoder"])
         self.image_encoder = VisualEncoder(**config["image_encoder"])
 
-        self._tokenizer = Tokenizer.from_file(tokenizer_path)
-        self._tokenizer.no_padding()
-        self._pad_token_idx = self.text_encoder.padding_idx
-
-        self._image_transform = Compose(
-            [
-                Resize(self._image_size, interpolation=InterpolationMode.BICUBIC),
-                convert_to_rgb,
-                CenterCrop(self._image_size),
-                ToTensor(),
-                Normalize(
-                    mean=(0.48145466, 0.4578275, 0.40821073),
-                    std=(0.26862954, 0.26130258, 0.27577711),
-                ),
-            ],
-        )
-
     def encode_image(
         self,
         images: Tensor,
@@ -426,28 +398,23 @@ def encode_multimodal(
         image_features: Optional[Tensor] = None,
         text_features: Optional[Tensor] = None,
         attention_mask: Optional[Tensor] = None,
-    ) -> Tensor:
+        return_scores: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
         """Passes preprocessed texts (or precomputed texts features) and
             preprocessed images (or precomputed images features) through multimodal encoded to produce multimodal joint embeddings.
 
         :param image: Preprocessed images
-        :param text: Preprocesses texts
+        :param text: Preprocessed texts
         :param image_features: Precomputed images features
         :param text_features: Precomputed text features
         :param attention_mask: Attention masks, not required if pass `text` instead of text_features
         """
 
-        assert (
-            image is not None or image_features is not None
-        ), "Either `image` or `image_features` should be non None"
-        assert (
-            text is not None or text_features is not None
-        ), "Either `text_data` or `text_features` should be non None"
+        assert image is not None or image_features is not None, "Either `image` or `image_features` should be non None"
+        assert text is not None or text_features is not None, "Either `text_data` or `text_features` should be non None"
 
         if text_features is not None:
-            assert (
-                attention_mask is not None
-            ), "if `text_features` is not None, then you should pass `attention_mask`"
+            assert attention_mask is not None, "if `text_features` is not None, then you should pass `attention_mask`"
 
         if image_features is None:
             image_features = self.image_encoder.forward_features(image)
@@ -458,11 +425,16 @@ def encode_multimodal(
                 text["attention_mask"],
             )
 
-        return self.text_encoder.forward_multimodal(
+        embeddings = self.text_encoder.forward_multimodal(
             text_features,
             attention_mask if attention_mask is not None else text["attention_mask"],
             image_features,
         )
+        
+        if return_scores:
+            return self.get_matching_scores(embeddings), embeddings
+
+        return embeddings
 
     def get_matching_scores(self, embeddings: Tensor) -> Tensor:
         """Computes the probability that there is a match between images and texts based on their multimodal embeddings
@@ -472,69 +444,20 @@ def get_matching_scores(self, embeddings: Tensor) -> Tensor:
 
         return self.text_encoder.forward_matching(embeddings)
 
-    def preprocess_text(self, texts: Union[str, List[str]]) -> Dict[str, Tensor]:
-        """Transforms one or more strings into dictionary with tokenized strings and attention masks.
-
-        :param texts: text of list of texts to tokenizer
-        """
-        if isinstance(texts, str):
-            texts = [texts]
-
-        input_ids = torch.full(
-            (len(texts), self.text_encoder.max_position_embeddings),
-            fill_value=self._pad_token_idx,
-            dtype=torch.int64,
-        )
-
-        attention_mask = torch.zeros(
-            len(texts),
-            self.text_encoder.max_position_embeddings,
-            dtype=torch.int32,
-        )
-        encoded = self._tokenizer.encode_batch(texts)
-
-        for i, seq in enumerate(encoded):
-            seq_len = min(len(seq), self.text_encoder.max_position_embeddings)
-            input_ids[i, :seq_len] = torch.LongTensor(
-                seq.ids[: self.text_encoder.max_position_embeddings],
-            )
-            attention_mask[i, :seq_len] = 1
-
-        return {"input_ids": input_ids, "attention_mask": attention_mask}
-
-    def preprocess_image(self, images: Union[Image, List[Image]]) -> Tensor:
-        """Transforms one or more Pillow images into Torch Tensors.
-
-        :param images: image or list of images to preprocess
-        """
-
-        if isinstance(images, list):
-            batch_images = torch.empty(
-                (len(images), 3, self._image_size, self._image_size),
-                dtype=torch.float32,
-            )
-
-            for i, image in enumerate(images):
-                batch_images[i] = self._image_transform(image)
-
-            return batch_images
-        else:
-            return self._image_transform(images).unsqueeze(0)
-
     def forward(
         self,
-        images: torch.Tensor,
-        texts: dict,
-    ):
+        images: Tensor,
+        texts: Dict[str, Tensor],
+    ) -> Union[Tensor, Tensor]:
         """Inference forward method
 
         :param images: Preprocessed images
         :param texts: Preprocessed texts
         :return: embeddings for images and texts
         """
-        _, embs_imgs = self.image_encoder(images)
-        _, embs_txts = self.text_encoder(texts)
-        return embs_imgs, embs_txts
+        _, image_embeddings = self.image_encoder(images)
+        _, text_embeddings = self.text_encoder(texts)
+        return image_embeddings, text_embeddings
 
     @property
     def text_features_dim(self) -> int:
@@ -558,168 +481,3 @@ def embedding_dim(self) -> int:
     def multimodal_embedding_dim(self) -> int:
         """Dimensionality of multimodal joint embedding."""
         return self.text_encoder.dim
-
-
-class TritonClient(VLM):
-    """
-    Nvidia Triton client to connect to the remote VLM inference server.
-    """
-
-    def __init__(self, tokenizer_path, pad_token_idx, url: str = "localhost:7001"):
-        import tritonclient.http as httpclient
-
-        self._client = httpclient
-        self._triton_client = self._client.InferenceServerClient(url=url)
-
-        self._tokenizer = Tokenizer.from_file(tokenizer_path)
-        self._tokenizer.no_padding()
-        self._pad_token_idx = pad_token_idx
-
-        self._image_transform = Compose(
-            [
-                Resize(self._image_size, interpolation=InterpolationMode.BICUBIC),
-                convert_to_rgb,
-                CenterCrop(self._image_size),
-                ToTensor(),
-                Normalize(
-                    mean=(0.48145466, 0.4578275, 0.40821073),
-                    std=(0.26862954, 0.26130258, 0.27577711),
-                ),
-            ],
-        )
-
-    def encode_image(
-        self,
-        imgs: Tensor,
-    ):
-        """
-        Passes the pre-processed images through `image_encoder` to produce image embeddings.
-
-        :param imgs: Preprocessed image
-        """
-
-        # images prep
-        inputs = []
-        outputs = []
-        imgs = imgs.cpu().detach().numpy()
-        inputs.append(self._client.InferInput("inputs", imgs.shape, "FP32"))
-        inputs[0].set_data_from_numpy(imgs)
-        outputs.append(self._client.InferRequestedOutput("output"))
-
-        # Querying the server
-        results = self._triton_client.infer(
-            model_name="vit",
-            inputs=inputs,
-            outputs=outputs,
-        )
-        output_data = torch.from_numpy(results.as_numpy("output"))
-        return output_data
-
-    def encode_text(
-        self,
-        text: Dict[str, Tensor],
-    ):
-        """
-        Passes the pre-processed texts through `text_encoder` to produce texts embeddings.
-
-        :param text: Dictionary with tokenized texts and attention masks
-        """
-
-        # texts prep
-        inputs = []
-        input_ids, attention_mask = text["input_ids"], text["attention_mask"]
-        input_ids = input_ids.type(dtype=torch.int32).cpu().detach().numpy()
-        attention_mask = attention_mask.type(dtype=torch.int32).cpu().detach().numpy()
-        inputs.append(
-            self._client.InferInput("attention_mask", attention_mask.shape, "INT32"),
-        )
-        inputs.append(self._client.InferInput("input_ids", input_ids.shape, "INT32"))
-        inputs[0].set_data_from_numpy(attention_mask)
-        inputs[1].set_data_from_numpy(input_ids)
-        test_output = self._client.InferRequestedOutput("output")
-
-        # Querying the server
-        results = self._triton_client.infer(
-            model_name="albef",
-            inputs=inputs,
-            outputs=[test_output],
-        )
-        output_vec = torch.from_numpy(results.as_numpy("output"))
-        return output_vec
-
-    def encode_multimodal(self, *args, **kwargs):
-        raise NotImplementedError("Multimodal encodings coming soon!")
-
-
-class VLM_IPU(VLM):
-    """
-    Code for GraphCore IPUs.
-    Please read User Guide if you want UForm to work on GraphCore hardware.
-    (https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/intro.html)
-    """
-
-    def __init__(self, config: Dict, tokenizer_path: PathLike):
-        import poptorch
-
-        self.poptorch = poptorch
-        super().__init__(config, tokenizer_path)
-
-    def recomputation_checkpoint(self, module):
-        """
-        Annotates the output of a module to be checkpointed instead of recomputed
-        """
-
-        def recompute_outputs(module, inputs, outputs):
-            if type(outputs) is tuple:
-                return tuple(self.poptorch.recomputationCheckpoint(y) for y in outputs)
-            else:
-                return self.poptorch.recomputationCheckpoint(outputs)
-
-        module.register_forward_hook(recompute_outputs)
-
-    def parallelize(self):
-        """
-        Splits the model layers between IPU devices.
-        """
-        print("---------- Device Allocation -----------")
-        print("image_encoder 0 ~ 6--> IPU 0")
-        for index in range(4):
-            layer = self.image_encoder.blocks[index]
-            self.recomputation_checkpoint(layer)
-            self.image_encoder.blocks[index] = self.poptorch.BeginBlock(
-                layer,
-                f"image_encoder_layer{index}",
-                ipu_id=0,
-            )
-
-        print("image_encoder 4 ~ 8 --> IPU 1")
-        for index in range(4, 8):
-            layer = self.image_encoder.blocks[index]
-            self.recomputation_checkpoint(layer)
-            self.image_encoder.blocks[index] = self.poptorch.BeginBlock(
-                layer,
-                f"image_encoder_layer{index}",
-                ipu_id=1,
-            )
-
-        print("image_encoder 8 ~ 12 --> IPU 2")
-        for index in range(8, 12):
-            layer = self.image_encoder.blocks[index]
-            self.recomputation_checkpoint(layer)
-            self.image_encoder.blocks[index] = self.poptorch.BeginBlock(
-                layer,
-                f"image_encoder_layer{index}",
-                ipu_id=2,
-            )
-
-        print("text_enocder 0 ~ 4 --> IPU 3")
-        for index in range(0, 4):
-            layer = self.text_encoder.blocks[index]
-            self.recomputation_checkpoint(layer)
-            self.text_encoder.blocks[index] = self.poptorch.BeginBlock(
-                layer,
-                f"text_encoder_layer{index}",
-                ipu_id=3,
-            )
-
-        return self
diff --git a/python/uform/torch_preprocessor.py b/python/uform/torch_preprocessor.py
new file mode 100644
index 0000000..8bdc70b
--- /dev/null
+++ b/python/uform/torch_preprocessor.py
@@ -0,0 +1,98 @@
+from os import PathLike
+from typing import Dict, List, Union
+
+import torch
+from PIL.Image import Image
+from tokenizers import Tokenizer
+from torch import Tensor
+from torchvision.transforms import (
+    CenterCrop,
+    Compose,
+    InterpolationMode,
+    Normalize,
+    Resize,
+    ToTensor,
+)
+
+
+# lambda is not pickable
+def convert_to_rgb(image):
+    return image.convert("RGB")
+
+
+class TorchProcessor:
+    def __init__(self, config: Dict, tokenizer_path: PathLike):
+        """
+        :param config: model config
+        :param tokenizer_path: path to tokenizer file
+        :param tensor_type: which tensors to return, either pt (PyTorch) or np (NumPy)
+        """
+
+        self._image_size = config["image_encoder"]["image_size"]
+        self._max_seq_len = config["text_encoder"]["max_position_embeddings"]
+        self._tokenizer = Tokenizer.from_file(tokenizer_path)
+        self._tokenizer.no_padding()
+        self._pad_token_idx = config["text_encoder"]["padding_idx"]
+
+        self._image_transform = Compose(
+            [
+                Resize(self._image_size, interpolation=InterpolationMode.BICUBIC),
+                convert_to_rgb,
+                CenterCrop(self._image_size),
+                ToTensor(),
+                Normalize(
+                    mean=(0.48145466, 0.4578275, 0.40821073),
+                    std=(0.26862954, 0.26130258, 0.27577711),
+                ),
+            ],
+        )
+
+    def preprocess_text(self, texts: Union[str, List[str]]) -> Dict[str, Tensor]:
+        """Transforms one or more strings into dictionary with tokenized strings and attention masks.
+
+        :param texts: text of list of texts to tokenizer
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+
+        input_ids = torch.full(
+            (len(texts), self._max_seq_len),
+            fill_value=self._pad_token_idx,
+            dtype=torch.int64,
+        )
+
+        attention_mask = torch.zeros(
+            len(texts),
+            self._max_seq_len,
+            dtype=torch.int32,
+        )
+        encoded = self._tokenizer.encode_batch(texts)
+
+        for i, seq in enumerate(encoded):
+            seq_len = min(len(seq), self._max_seq_len)
+            input_ids[i, :seq_len] = torch.LongTensor(
+                seq.ids[:seq_len],
+            )
+            attention_mask[i, :seq_len] = 1
+
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+    def preprocess_image(self, images: Union[Image, List[Image]]) -> Tensor:
+        """Transforms one or more Pillow images into Torch Tensors.
+
+        :param images: image or list of images to preprocess
+        """
+
+        if isinstance(images, list):
+            batch_images = torch.empty(
+                (len(images), 3, self._image_size, self._image_size),
+                dtype=torch.float32,
+            )
+
+            for i, image in enumerate(images):
+                batch_images[i] = self._image_transform(image)
+
+        else:
+            batch_images = self._image_transform(images).unsqueeze(0)
+
+        return batch_images
diff --git a/src/uform/__init__.py b/src/uform/__init__.py
deleted file mode 100755
index f02bd6c..0000000
--- a/src/uform/__init__.py
+++ /dev/null
@@ -1,82 +0,0 @@
-from json import load
-from typing import Mapping, Optional, Tuple
-
-import torch
-from huggingface_hub import snapshot_download
-
-from uform.models import (
-    MLP,
-    VLM,
-    VLM_IPU,
-    Attention,
-    LayerScale,
-    TextEncoder,
-    TextEncoderBlock,
-    TritonClient,
-    VisualEncoder,
-    VisualEncoderBlock,
-    convert_to_rgb,
-)
-
-__all__ = [
-    "MLP",
-    "VLM",
-    "VLM_IPU",
-    "Attention",
-    "LayerScale",
-    "TextEncoder",
-    "TextEncoderBlock",
-    "TritonClient",
-    "VisualEncoder",
-    "VisualEncoderBlock",
-    "convert_to_rgb",
-    "get_checkpoint",
-    "get_model",
-    "get_client",
-    "get_model_ipu",
-]
-
-
-def get_checkpoint(model_name, token) -> Tuple[str, Mapping, str]:
-    model_path = snapshot_download(repo_id=model_name, token=token)
-    config_path = f"{model_path}/torch_config.json"
-    state = torch.load(f"{model_path}/torch_weight.pt")
-
-    return config_path, state, f"{model_path}/tokenizer.json"
-
-
-def get_model(model_name: str, token: Optional[str] = None) -> VLM:
-    config_path, state, tokenizer_path = get_checkpoint(model_name, token)
-
-    with open(config_path) as f:
-        model = VLM(load(f), tokenizer_path)
-
-    model.image_encoder.load_state_dict(state["image_encoder"])
-    model.text_encoder.load_state_dict(state["text_encoder"])
-
-    return model.eval()
-
-
-def get_client(
-    url: str,
-    model_name: str = "unum-cloud/uform-vl-english",
-    token: Optional[str] = None,
-) -> TritonClient:
-    config_path, _, tokenizer_path = get_checkpoint(model_name, token)
-
-    with open(config_path) as f:
-        pad_token_idx = load(f)["text_encoder"]["padding_idx"]
-
-    return TritonClient(tokenizer_path, pad_token_idx, url)
-
-
-def get_model_ipu(model_name: str, token: Optional[str] = None) -> VLM_IPU:
-    config_path, state, tokenizer_path = get_checkpoint(model_name, token)
-
-    with open(config_path) as f:
-        model = VLM_IPU(load(f), tokenizer_path)
-
-    model.image_encoder.load_state_dict(state["image_encoder"])
-    model.text_encoder.load_state_dict(state["text_encoder"])
-
-    return model.eval()


Test patch:
diff --git a/python/scripts/test_embeddings.py b/python/scripts/test_embeddings.py
new file mode 100644
index 0000000..f831d72
--- /dev/null
+++ b/python/scripts/test_embeddings.py
@@ -0,0 +1,143 @@
+from typing import Tuple
+
+import pytest
+from PIL import Image
+import uform
+
+# PyTorch is a very heavy dependency, so we may want to skip these tests if it's not installed
+try:
+    import torch
+
+    torch_available = True
+except:
+    torch_available = False
+    
+# ONNX is not a very light dependency either
+try:
+    import onnx
+
+    onnx_available = True
+except:
+    onnx_available = False
+
+torch_models = [
+    "unum-cloud/uform-vl-english",
+    "unum-cloud/uform-vl-multilingual-v2",
+]
+
+onnx_models_and_providers = [
+    ("unum-cloud/uform-vl-english-small", "cpu", "fp32"),
+    ("unum-cloud/uform-vl-english-large", "cpu", "fp32"),
+    ("unum-cloud/uform-vl-english-small", "gpu", "fp32"),
+    ("unum-cloud/uform-vl-english-large", "gpu", "fp32"),
+    ("unum-cloud/uform-vl-english-small", "gpu", "fp16"),
+    ("unum-cloud/uform-vl-english-large", "gpu", "fp16"),
+]
+
+@pytest.mark.skipif(not torch_available, reason="PyTorch is not installed")
+@pytest.mark.parametrize("model_name", torch_models)
+def test_torch_one_embedding(model_name: str):
+    model, processor = uform.get_model(model_name)
+    text = "a small red panda in a zoo"
+    image_path = "assets/unum.png"
+
+    image = Image.open(image_path)
+    image_data = processor.preprocess_image(image)
+    text_data = processor.preprocess_text(text)
+
+    image_features, image_embedding = model.encode_image(image_data, return_features=True)
+    text_features, text_embedding = model.encode_text(text_data, return_features=True)
+
+    assert image_embedding.shape[0] == 1, "Image embedding batch size is not 1"
+    assert text_embedding.shape[0] == 1, "Text embedding batch size is not 1"
+
+    # Test reranking
+    score, joint_embedding = model.encode_multimodal(
+        image_features=image_features,
+        text_features=text_features,
+        attention_mask=text_data["attention_mask"],
+        return_scores=True,
+    )
+    assert score.shape[0] == 1, "Matching score batch size is not 1"
+    assert joint_embedding.shape[0] == 1, "Joint embedding batch size is not 1"
+
+
+@pytest.mark.skipif(not torch_available, reason="PyTorch is not installed")
+@pytest.mark.parametrize("model_name", torch_models)
+@pytest.mark.parametrize("batch_size", [1, 2])
+def test_torch_many_embeddings(model_name: str, batch_size: int):
+    model, processor = uform.get_model(model_name)
+    texts = ["a small red panda in a zoo"] * batch_size
+    image_paths = ["assets/unum.png"] * batch_size
+
+    images = [Image.open(path) for path in image_paths]
+    image_data = processor.preprocess_image(images)
+    text_data = processor.preprocess_text(texts)
+
+    image_embeddings = model.encode_image(image_data, return_features=False)
+    text_embeddings = model.encode_text(text_data, return_features=False)
+
+    assert image_embeddings.shape[0] == batch_size, "Image embedding is unexpected"
+    assert text_embeddings.shape[0] == batch_size, "Text embedding is unexpected"
+
+
+@pytest.mark.skipif(not onnx_available, reason="ONNX is not installed")
+@pytest.mark.parametrize("model_specs", onnx_models_and_providers)
+def test_onnx_one_embedding(model_specs: Tuple[str, str, str]):
+
+    from uform.onnx_models import ExecutionProviderError
+
+    try:
+
+        model, processor = uform.get_model_onnx(*model_specs)
+        text = "a small red panda in a zoo"
+        image_path = "assets/unum.png"
+
+        image = Image.open(image_path)
+        image_data = processor.preprocess_image(image)
+        text_data = processor.preprocess_text(text)
+
+        image_features, image_embedding = model.encode_image(image_data, return_features=True)
+        text_features, text_embedding = model.encode_text(text_data, return_features=True)
+
+        assert image_embedding.shape[0] == 1, "Image embedding batch size is not 1"
+        assert text_embedding.shape[0] == 1, "Text embedding batch size is not 1"
+
+        score, joint_embedding = model.encode_multimodal(
+            image_features=image_features,
+            text_features=text_features,
+            attention_mask=text_data["attention_mask"],
+            return_scores=True,
+        )
+        assert score.shape[0] == 1, "Matching score batch size is not 1"
+        assert joint_embedding.shape[0] == 1, "Joint embedding batch size is not 1"
+
+    except ExecutionProviderError as e:
+        pytest.skip(f"Execution provider error: {e}")
+
+
+@pytest.mark.skipif(not onnx_available, reason="ONNX is not installed")
+@pytest.mark.parametrize("model_specs", onnx_models_and_providers)
+@pytest.mark.parametrize("batch_size", [1, 2])
+def test_onnx_many_embeddings(model_specs: Tuple[str, str, str], batch_size: int):
+
+    from uform.onnx_models import ExecutionProviderError
+
+    try:
+
+        model, processor = uform.get_model_onnx(*model_specs)
+        texts = ["a small red panda in a zoo"] * batch_size
+        image_paths = ["assets/unum.png"] * batch_size
+
+        images = [Image.open(path) for path in image_paths]
+        image_data = processor.preprocess_image(images)
+        text_data = processor.preprocess_text(texts)
+
+        image_embeddings = model.encode_image(image_data, return_features=False)
+        text_embeddings = model.encode_text(text_data, return_features=False)
+
+        assert image_embeddings.shape[0] == batch_size, "Image embedding is unexpected"
+        assert text_embeddings.shape[0] == batch_size, "Text embedding is unexpected"
+
+    except ExecutionProviderError as e:
+        pytest.skip(f"Execution provider error: {e}")
diff --git a/python/scripts/test_generative.py b/python/scripts/test_generative.py
new file mode 100644
index 0000000..0d38b64
--- /dev/null
+++ b/python/scripts/test_generative.py
@@ -0,0 +1,73 @@
+import pytest
+from PIL import Image
+
+# PyTorch is a very heavy dependency, so we may want to skip these tests if it's not installed
+try:
+    import torch
+
+    torch_available = True
+except:
+    torch_available = False
+
+torch_hf_models = [
+    "unum-cloud/uform-gen2-qwen-500m",
+    "unum-cloud/uform-gen2-dpo",
+]
+
+
+@pytest.mark.skipif(not torch_available, reason="PyTorch is not installed")
+@pytest.mark.parametrize("model_name", torch_hf_models)
+def test_one_conversation(model_name: str):
+    from transformers import AutoModel, AutoProcessor
+
+    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+
+    prompt = "Describe the image in great detail."
+    image = Image.open("assets/unum.png")
+
+    inputs = processor(text=[prompt], images=[image], return_tensors="pt")
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            do_sample=False,
+            use_cache=True,
+            max_new_tokens=10,
+            pad_token_id=processor.tokenizer.pad_token_id,
+        )
+    prompt_len = inputs["input_ids"].shape[1]
+    decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
+
+    assert len(decoded_text), "No text was generated from the model."
+
+
+@pytest.mark.skipif(not torch_available, reason="PyTorch is not installed")
+@pytest.mark.parametrize("model_name", torch_hf_models)
+@pytest.mark.parametrize("batch_size", [1, 2])
+def test_many_conversations(model_name: str, batch_size: int):
+
+    from transformers import AutoModel, AutoProcessor
+
+    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+
+    prompt = "Describe the image in great detail."
+    image = Image.open("assets/unum.png")
+
+    texts = [prompt] * batch_size
+    images = [image] * batch_size
+    inputs = processor(text=texts, images=images, return_tensors="pt")
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            do_sample=False,
+            use_cache=True,
+            max_new_tokens=10,
+            pad_token_id=processor.tokenizer.pad_token_id,
+        )
+    prompt_len = inputs["input_ids"].shape[1]
+    decoded_texts = processor.batch_decode(output[:, prompt_len:])
+
+    assert all(len(decoded_text) for decoded_text in decoded_texts), "No text was generated from the model."