From aefb710d01b02e0914af29246d2903d97265c1d3 Mon Sep 17 00:00:00 2001
From: michaelkillgta <269775689+michaelkillgta@users.noreply.github.com>
Date: Mon, 12 Oct 2026 19:55:58 +0000
Subject: [PATCH] fix(m3): score a single pair on multiple devices

---
 .../inference/embedder/encoder_only/m3.py     |  7 +++
 tests/test_m3_compute_score_multi_device.py   | 49 +++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 tests/test_m3_compute_score_multi_device.py

diff --git a/FlagEmbedding/inference/embedder/encoder_only/m3.py b/FlagEmbedding/inference/embedder/encoder_only/m3.py
index dc533f55..50ceda04 100644
--- a/FlagEmbedding/inference/embedder/encoder_only/m3.py
+++ b/FlagEmbedding/inference/embedder/encoder_only/m3.py
@@ -527,6 +527,11 @@ def compute_score(
                 **kwargs
             )
 
+        # A single (query, passage) pair is scored as a one-pair batch, then unwrapped like on a single device
+        one_input_pair = isinstance(sentence_pairs[0], str)
+        if one_input_pair:
+            sentence_pairs = [sentence_pairs]
+
         pool = self.start_multi_process_pool(M3Embedder._compute_score_multi_process_worker)
         embeddings = self.compute_score_multi_process(
             sentence_pairs,
@@ -538,6 +543,8 @@ def compute_score(
             **kwargs
         )
         self.stop_multi_process_pool(pool)
+        if one_input_pair:
+            return {k: v[0] for k, v in embeddings.items()}
         return embeddings
 
     # adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L877
diff --git a/tests/test_m3_compute_score_multi_device.py b/tests/test_m3_compute_score_multi_device.py
new file mode 100644
index 00000000..9ed751a2
--- /dev/null
+++ b/tests/test_m3_compute_score_multi_device.py
@@ -0,0 +1,49 @@
+"""
+Test that BGE-M3 compute_score handles a single (query, passage) pair on multiple devices.
+
+The multi-process pool is stubbed, so no model is loaded and no processes are started.
+"""
+from FlagEmbedding.inference.embedder.encoder_only.m3 import M3Embedder
+
+MODES = ["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"]
+
+
+def _make_embedder(received):
+    embedder = M3Embedder.__new__(M3Embedder)
+    embedder.target_devices = ["cpu", "cpu"]
+    embedder.batch_size = 256
+    embedder.query_max_length = 512
+    embedder.passage_max_length = 512
+    embedder.start_multi_process_pool = lambda worker: {"processes": [None, None]}
+    embedder.stop_multi_process_pool = lambda pool: None
+
+    def compute_score_multi_process(sentence_pairs, pool, **kwargs):
+        received.append(sentence_pairs)
+        return {mode: [0.5] * len(sentence_pairs) for mode in MODES}
+
+    embedder.compute_score_multi_process = compute_score_multi_process
+    return embedder
+
+
+def test_m3_compute_score_single_pair_multi_device():
+    received = []
+    embedder = _make_embedder(received)
+
+    pair = ("what is BGE M3?", "BGE M3 is an embedding model.")
+    scores = embedder.compute_score(pair)
+
+    # The pair is scored as one pair, not as two separate strings
+    assert received == [[pair]]
+    # Same output shape as a single device: one float per mode
+    assert scores == {mode: 0.5 for mode in MODES}
+
+
+def test_m3_compute_score_pair_list_multi_device():
+    received = []
+    embedder = _make_embedder(received)
+
+    pairs = [("q1", "p1"), ("q2", "p2")]
+    scores = embedder.compute_score(pairs)
+
+    assert received == [pairs]
+    assert scores == {mode: [0.5, 0.5] for mode in MODES}