1from sentence_transformers import CrossEncoder
2import torch
3
4MODEL_NAME = "hotchpotch/japanese-reranker-small-v2"
5
6model = CrossEncoder(MODEL_NAME)
7if model.device == "cuda" or model.device == "mps":
8 model.model.half()
9query = "感動的な映画について"
10passages = [
11 "深いテーマを持ちながらも、観る人の心を揺さぶる名作。登場人物の心情描写が秀逸で、ラストは涙なしでは見られない。",
12 "重要なメッセージ性は評価できるが、暗い話が続くので気分が落ち込んでしまった。もう少し明るい要素があればよかった。",
13 "どうにもリアリティに欠ける展開が気になった。もっと深みのある人間ドラマが見たかった。",
14 "アクションシーンが楽しすぎる。見ていて飽きない。ストーリーはシンプルだが、それが逆に良い。",
15]
16scores = model.predict(
17 [(query, passage) for passage in passages],
18 show_progress_bar=True,
19)
20print("Scores:", scores)
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2from torch.nn import Sigmoid
3
4
5def detect_device():
6 if torch.cuda.is_available():
7 return "cuda"
8 elif hasattr(torch, "mps") and torch.mps.is_available():
9 return "mps"
10 return "cpu"
11
12
13device = detect_device()
14tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
15model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
16model.to(device)
17model.eval()
18
19if device == "cuda":
20 model.half()
21
22query = "感動的な映画について"
23passages = [
24 "深いテーマを持ちながらも、観る人の心を揺さぶる名作。登場人物の心情描写が秀逸で、ラストは涙なしでは見られない。",
25 "重要なメッセージ性は評価できるが、暗い話が続くので気分が落ち込んでしまった。もう少し明るい要素があればよかった。",
26 "どうにもリアリティに欠ける展開が気になった。もっと深みのある人間ドラマが見たかった。",
27 "アクションシーンが楽しすぎる。見ていて飽きない。ストーリーはシンプルだが、それが逆に良い。",
28]
29inputs = tokenizer(
30 [(query, passage) for passage in passages],
31 padding=True,
32 truncation=True,
33 max_length=512,
34 return_tensors="pt",
35)
36inputs = {k: v.to(device) for k, v in inputs.items()}
37logits = model(**inputs).logits
38activation = Sigmoid()
39scores = activation(logits).squeeze().tolist()
40
41print("Scores:", scores)
japanese-reranker-tiny-v2、japanese-reranker-xsmall-v2、japanese-reranker-small-v2、japanese-reranker-base-v2は、以下の特徴を持つ小型リランカーモデルです:
以下は約15万ペアをリランキングした際の推論速度結果(トークナイズ時間を除く純粋なモデル推論時間)です。MPS(Appleシリコン)とCPU測定にはM4 Max、GPUにはRTX5090を使用しています。GPU処理では flash-attention2 を使用しています。