Gemma 4 E2B It Assistant ONNX

--- library_name: transformers base_model: google/gemma-4-E2B-it license: apache-2.0 license_link: https://ai.google.dev/gemma/docs/gemma_4_license pipeline_tag: any-to-any --- # Gemma 4 E2B IT Assistant ONNX This repository contains an ONNX conversion of the original [`google/gemma-4-E2B-it-assistant`](https://huggingface.co/google/gemma-4-E2B-it-assistant) model for use with ONNX runtimes and transformers.js. The original model card is included below. --- Hugging Face | GitHub | MTP Documentation License : Apache 2.0 | Authors : Google DeepMind > [!Note] > This model card is for the Multi-Token Prediction (MTP) drafters for the Gemma 4 models. MTP is implemented by extending the base model with a smaller, faster draft model. When used in a Speculative Decoding pipeline, the draft model predicts several tokens ahead, which the target model then verifies in parallel. This results in significant decoding speedups (up to 3x) while guaranteeing the exact same quality as standard generation, making these checkpoints perfect for low-latency and on-device applications. Gemma is a family of open models built by Google DeepMind. Gemma 4 models are multimodal, handling text and image input (with audio supported on small models) and generating text output. This release includes open-weights models in both pre-trained and instruction-tuned variants. Gemma 4 features a context window of up to 256K tokens and maintains multilingual support in over 140 languages. Featuring both Dense and Mixture-of-Experts (MoE) architectures, Gemma 4 is well-suited for tasks like text generation, coding, and reasoning. The models are available in four distinct sizes: **E2B**, **E4B**, **26B A4B**, and **31B**. Their diverse sizes make them deployable in environments ranging from high-end phones to laptops and servers, democratizing access to state-of-the-art AI. Gemma 4 introduces key **capability and architectural advancements**: * **Reasoning** – All models in the family are designed as highly capable reasoners, with configurable thinking modes. * **Extended Multimodalities** – Processes Text, Image with variable aspect ratio and resolution support (all models), Video, and Audio (featured natively on the E2B and E4B models). * **Diverse & Efficient Architectures** – Offers Dense and Mixture-of-Experts (MoE) variants of different sizes for scalable deployment. * **Optimized for On-Device** – Smaller models are specifically designed for efficient local execution on laptops and mobile devices. * **Increased Context Window** – The small models feature a 128K context window, while the medium models support 256K. * **Enhanced Coding & Agentic Capabilities** – Achieves notable improvements in coding benchmarks alongside native function-calling support, powering highly capable autonomous agents. * **Native System Prompt Support** – Gemma 4 introduces native support for the `system` role, enabling more structured and controllable conversations. ## **Models Overview** Gemma 4 models are designed to deliver frontier-level performance at each size, targeting deployment scenarios from mobile and edge devices (E2B, E4B) to consumer GPUs and workstations (26B A4B, 31B). They are well-suited for reasoning, agentic workflows, coding, and multimodal understanding. The models employ a hybrid attention mechanism that interleaves local sliding window attention with full global attention, ensuring the final layer is always global. This hybrid design delivers the processing speed and low memory footprint of a lightweight model without sacrificing the deep awareness required for complex, long-context tasks. To optimize memory for long contexts, global layers feature unified Keys and Values, and apply Proportional RoPE (p-RoPE). ### Dense Models | Property | E2B | E4B | 31B Dense | | :---- | :---- | :---- | :---- | | **Total Parameters** | 2.3B effective (5.1B with embeddings) | 4.5B effective (8B with embeddings) | 30.7B | | **Layers** | 35 | 42 | 60 | | **Sliding Window** | 512 tokens | 512 tokens | 1024 tokens | | **Context Length** | 128K tokens | 128K tokens | 256K tokens | | **Vocabulary Size** | 262K | 262K | 262K | | **Supported Modalities** | Text, Image, Audio | Text, Image, Audio | Text, Image | | **Vision Encoder Parameters** | *~150M* | *~150M* | *~550M* | | **Audio Encoder Parameters** | *~300M* | *~300M* | No Audio | The "E" in E2B and E4B stands for "effective" parameters. The smaller models incorporate Per-Layer Embeddings (PLE) to maximize parameter efficiency in on-device deployments. Rather than adding more layers or parameters to the model, PLE gives each decoder layer its own small embedding for every token. These embedding tables are large but are only used for quick lookups, which is why the effective parameter count is much smaller than the total. ### Mixture-of-Experts (MoE) Model | Property | 26B A4B MoE | | :---- | :---- | | **Total Parameters** | 25.2B | | **Active Parameters** | 3.8B | | **Layers** | 30 | | **Sliding Window** | 1024 tokens | | **Context Length** | 256K tokens | | **Vocabulary Size** | 262K | | **Expert Count** | 8 active / 128 total and 1 shared | | **Supported Modalities** | Text, Image | | **Vision Encoder Parameters** | *~550M* | The "A" in 26B A4B stands for "active parameters" in contrast to the total number of parameters the model contains. By only activating a 4B subset of parameters during inference, the Mixture-of-Experts model runs much faster than its 26B total might suggest. This makes it an excellent choice for fast inference compared to the dense 31B model since it runs almost as fast as a 4B-parameter model. ## **Benchmark Results** These models were evaluated against a large collection of different datasets and metrics to cover different aspects of text generation. Evaluation results marked in the table are for instruction-tuned models. | | Gemma 4 31B | Gemma 4 26B A4B | Gemma 4 E4B | Gemma 4 E2B | Gemma 3 27B (no think) | | :---- | :---- | :---- | :---- | :---- | :---- | | MMLU Pro | 85.2% | 82.6% | 69.4% | 60.0% | 67.6% | | AIME 2026 no tools | 89.2% | 88.3% | 42.5% | 37.5% | 20.8% | | LiveCodeBench v6 | 80.0% | 77.1% | 52.0% | 44.0% | 29.1% | | Codeforces ELO | 2150 | 1718 | 940 | 633 | 110 | | GPQA Diamond | 84.3% | 82.3% | 58.6% | 43.4% | 42.4% | | Tau2 (average over 3) | 76.9% | 68.2% | 42.2% | 24.5% | 16.2% | | HLE no tools | 19.5% | 8.7% | - | - | - | | HLE with search | 26.5% | 17.2% | - | - | - | | BigBench Extra Hard | 74.4% | 64.8% | 33.1% | 21.9% | 19.3% | | MMMLU | 88.4% | 86.3% | 76.6% | 67.4% | 70.7% | | **Vision** | | | | | | | MMMU Pro | 76.9% | 73.8% | 52.6% | 44.2% | 49.7% | | OmniDocBench 1.5 (average edit distance, lower is better) | 0.131 | 0.149 | 0.181 | 0.290 | 0.365 | | MATH-Vision | 85.6% | 82.4% | 59.5% | 52.4% | 46.0% | | MedXPertQA MM | 61.3% | 58.1% | 28.7% | 23.5% | - | | **Audio** | | | | | | | CoVoST | - | - | 35.54 | 33.47 | - | | FLEURS (lower is better) | - | - | 0.08 | 0.09 | - | | **Long Context** | | | | | | | MRCR v2 8 needle 128k (average) | 66.4% | 44.1% | 25.4% | 19.1% | 13.5% | ## **Core Capabilities** Gemma 4 models handle a broad range of tasks across text, vision, and audio. Key capabilities include: * **Thinking** – Built-in reasoning mode that lets the model think step-by-step before answering. * **Long Context** – Context windows of up to 128K tokens (E2B/E4B) and 256K tokens (26B A4B/31B). * **Image Understanding** – Object detection, Document/PDF parsing, screen and UI understanding, chart comprehension, OCR (including multilingual), handwriting recognition, and pointing. Images can be processed at variable aspect ratios and resolutions. * **Video Understanding** – Analyze video by processing sequences of frames. * **Interleaved Multimodal Input** – Freely mix text and images in any order within a single prompt. * **Function Calling** – Native support for structured tool use, enabling agentic workflows. * **Coding** – Code generation, completion, and correction. * **Multilingual** – Out-of-the-box support for 35+ languages, pre-trained on 140+ languages. * **Audio** (E2B and E4B only) – Automatic speech recognition (ASR) and speech-to-translated-text translation across multiple languages. ## Getting Started You can use all Gemma 4 models with the latest version of Transformers. To get started, install the necessary dependencies in your environment: `pip install -U transformers torch accelerate` Once you have everything installed, you can proceed to load the target model and the assistant model with the code below: ```python from transformers import AutoProcessor, AutoModelForCausalLM TARGET_MODEL_ID = "google/gemma-4-E2B-it" ASSISTANT_MODEL_ID = "google/gemma-4-E2B-it-assistant" # Target Model processor = AutoProcessor.from_pretrained(TARGET_MODEL_ID) target_model = AutoModelForCausalLM.from_pretrained( TARGET_MODEL_ID, dtype="auto", device_map="auto", ) # Assistant Model (the drafter) assistant_model = AutoModelForCausalLM.from_pretrained( ASSISTANT_MODEL_ID, dtype="auto", device_map="auto", ) ``` Once the models are loaded, you can start generating output: ```python # Prompt messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Write a short joke about saving RAM."}, ] # Process input text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) inputs = processor(text=text, return_tensors="pt").to(target_model.device) input_len = inputs["input_ids"].shape[-1] # Generate output outputs = target_model.generate( **inputs, assistant_model=assistant_model, max_new_tokens=256, ) response = processor.decode(outputs[0][input_len:], skip_special_tokens=False) # Parse output processor.parse_response(response) ``` To enable reasoning, set `enable_thinking=True` and the `parse_response` function will take care of parsing the thinking output. Below, you will also find snippets for processing audio (E2B and E4B only), images, and video alongside text: Code for processing Audio Instead of using `AutoModelForCausalLM`, you can use `AutoModelForMultimodalLM` to process audio. To use it, make sure to install the following packages: `pip install -U transformers torch torchvision librosa accelerate` Once you have everything installed, you can proceed to load the target model and the assistant model with the code below: ```python import torch from transformers import AutoProcessor, AutoModelForCausalLM, AutoModelForMultimodalLM TARGET_MODEL_ID = "google/gemma-4-E2B-it" ASSISTANT_MODEL_ID = "google/gemma-4-E2B-it-assistant" # Target Model processor = AutoProcessor.from_pretrained(TARGET_MODEL_ID) target_model = AutoModelForMultimodalLM.from_pretrained( TARGET_MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto", ) # Assistant Model (the drafter) assistant_model = AutoModelForCausalLM.from_pretrained( ASSISTANT_MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto", ) ``` Once the model is loaded, you can start generating output by directly referencing the audio URL in the prompt: ```python # Prompt - add audio after text messages = [ { "role": "user", "content": [ {"type": "text", "text": "Transcribe the following speech segment in its original language. Follow these specific instructions for formatting the answer:\n* Only output the transcription, with no newlines.\n* When transcribing numbers, write the digits, i.e. write 1.7 and not one point seven, and write 3 instead of three."}, {"type": "audio", "audio": "https://github.com/google-gemma/cookbook/raw/refs/heads/main/apps/sample-data/journal1.wav"}, ] } ] # Process input text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) inputs = processor(text=text, return_tensors="pt").to(target_model.device) input_len = inputs["input_ids"].shape[-1] # Generate output outputs = target_model.generate( **inputs, assistant_model=assistant_model, max_new_tokens=256, ) response =...