japanese-reranker-base-v2-onnx

kens24481/japanese-reranker-base-v2-onnx

A Japanese reranker model: an ONNX (fp32) export of hotchpotch/japanese-reranker-base-v2. It can be used without PyTorch, via onnxruntime or fastembed.

The weights are unchanged (no quantization, pruning, or re-training); only the file format has been converted. For model performance, evaluation results, and training data, see the original model card.

Files

ONNX model inputs and outputs

• Inputs: input_ids, attention_mask (both int64, shape [batch_size, sequence_length])
• Output: logits (float32, shape [batch_size, 1]). Apply a sigmoid to obtain a relevance score in [0, 1]

Usage

fastembed

from fastembed.common.model_description import ModelSource
from fastembed.rerank.cross_encoder import TextCrossEncoder

MODEL_NAME = "kens24481/japanese-reranker-base-v2-onnx"

TextCrossEncoder.add_custom_model(
    model=MODEL_NAME,
    sources=ModelSource(hf=MODEL_NAME),
    model_file="onnx/model.onnx",
)
encoder = TextCrossEncoder(MODEL_NAME)

query = "感動的な映画について"
passages = [
    "深いテーマを持ちながらも、観る人の心を揺さぶる名作。登場人物の心情描写が秀逸で、ラストは涙なしでは見られない。",
    "重要なメッセージ性は評価できるが、暗い話が続くので気分が落ち込んでしまった。もう少し明るい要素があればよかった。",
    "どうにもリアリティに欠ける展開が気になった。もっと深みのある人間ドラマが見たかった。",
    "アクションシーンが楽しすぎる。見ていて飽きない。ストーリーはシンプルだが、それが逆に良い。",
]
scores = list(encoder.rerank(query, passages))  # logits (before sigmoid)
print(scores)

onnxruntime

import numpy as np
import onnxruntime as ort
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer

model_dir = snapshot_download("kens24481/japanese-reranker-base-v2-onnx")
tokenizer = Tokenizer.from_file(f"{model_dir}/tokenizer.json")
tokenizer.enable_truncation(max_length=512)
tokenizer.enable_padding()
session = ort.InferenceSession(
    f"{model_dir}/onnx/model.onnx", providers=["CPUExecutionProvider"]
)

query = "感動的な映画について"
passages = ["深いテーマを持ちながらも、観る人の心を揺さぶる名作。", "アクションシーンが楽しすぎる。"]
encodings = tokenizer.encode_batch([(query, p) for p in passages])
feed = {
    "input_ids": np.array([e.ids for e in encodings], dtype=np.int64),
    "attention_mask": np.array([e.attention_mask for e in encodings], dtype=np.int64),
}
logits = session.run(None, feed)[0].squeeze(-1)
scores = 1 / (1 + np.exp(-logits))
print(scores)

Conversion and verification

The model was converted with convert.py (Python 3.11, torch 2.2.2, transformers 4.57.6, onnx 1.23.0, onnxruntime 1.23.2).

python convert.py --output-dir .

On the verification inputs, the maximum difference in scores (after sigmoid) between the original model (PyTorch) and this model (onnxruntime, CPU) was 4.5e-08.

License

This model is released under the MIT License (LICENSE). The licenses of the original model and its base models are listed below. See NOTICE for details.

Acknowledgements

Thanks to hotchpotch for publishing the original model, and to Nagoya University (cl-nagoya) and SB Intuitions for publishing the base models.