naranor/ModernBERT-large-ONNX-Attentions

ModernBERT-large ONNX with Attention Weights (INT8)

This repository provides a quantized INT8 ONNX export of answerdotai/ModernBERT-large. This version is specifically modified to expose raw attention weights for each layer, enabling advanced context analysis and pruning techniques.

Key Features

• Architecture: Modernized BERT with Rotary Positional Embeddings (RoPE) and GeGLU.
• Context Window: Native support for up to 8,192 tokens.
• Modification: Custom ONNX graph export that includes attentions.{0..27} as additional outputs.
• Optimization: INT8 quantization (AVX-512 VNNI optimized) for high-performance CPU inference.

Usage with WAMP

This model is a primary component of the Weighted Attention Message Pruner (WAMP-proxy). WAMP uses these attention weights to identify and retain the most important messages in long chat histories, saving up to 50% on LLM token costs.

Quick Inference (ONNX Runtime)

import onnxruntime as ort
from transformers import AutoTokenizer
import numpy as np

# Load model and tokenizer
session = ort.InferenceSession("model_quantized.onnx")
tokenizer = AutoTokenizer.from_pretrained("naranor/ModernBERT-large-ONNX-Attentions")

# Prepare input
inputs = tokenizer("Your text here", return_tensors="np")
onnx_inputs = {
    "input_ids": inputs["input_ids"],
    "attention_mask": inputs["attention_mask"]
}

# Run with attentions
outputs = session.run(None, onnx_inputs)
# Last hidden state is outputs[0], attentions are in subsequent indices

Attribution & Original Work

• Original Model: ModernBERT-large by AnswerDotAI and LightOn.
• Exported by: naranor using the WAMP Universal Exporter.

License

This model is licensed under the Apache License 2.0. You are free to use, modify, and distribute this model, including for commercial purposes, as long as you provide attribution to the original authors.

---

License details: Apache 2.0