multi-GPU support for mine_hard_negatives

Added support for multi-GPU encoding in sentence embeddings with model.encode_multi_process

sentence_transformers/util.py

@@ -533,6 +533,7 @@ def mine_hard_negatives(
     faiss_batch_size: int = 16384,
     use_faiss: bool = False,
     verbose: bool = True,
+    use_multiple_gpus = True,
 ) -> Dataset:
     """
     Add hard negatives to a dataset of (anchor, positive) pairs to create (anchor, positive, negative) triplets or
@@ -714,12 +715,21 @@ def mine_hard_negatives(
         except Exception:
             pass
 
-        corpus_embeddings = model.encode(
-            corpus, batch_size=batch_size, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True
-        )
-        query_embeddings = model.encode(
-            queries, batch_size=batch_size, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True
-        )
+        if use_multiple_gpus:
+            pool = model.start_multi_process_pool()
+
+            corpus_embeddings = model.encode_multi_process(
+                corpus, pool, batch_size=batch_size, normalize_embeddings=True, show_progress_bar=True)
+
+            query_embeddings = model.encode_multi_process(
+                queries, pool, batch_size=batch_size, normalize_embeddings=True, show_progress_bar=True)
+        else
+            corpus_embeddings = model.encode(
+                corpus, batch_size=batch_size, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True
+            )
+            query_embeddings = model.encode(
+                queries, batch_size=batch_size, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True
+            )
         index.add(corpus_embeddings)
 
         scores_list = []
@@ -735,12 +745,22 @@ def mine_hard_negatives(
 
     else:
         # Embed the corpus and the queries
-        corpus_embeddings = model.encode(
-            corpus, batch_size=batch_size, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True
-        )
-        query_embeddings = model.encode(
-            queries, batch_size=batch_size, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True
-        )
+
+        if use_multiple_gpus:
+            pool = model.start_multi_process_pool()
+
+            corpus_embeddings = model.encode_multi_process(
+                corpus, pool, batch_size=batch_size, normalize_embeddings=True, show_progress_bar=True)
+
+            query_embeddings = model.encode_multi_process(
+                queries, pool, batch_size=batch_size, normalize_embeddings=True, show_progress_bar=True)
+        else
+            corpus_embeddings = model.encode(
+                corpus, batch_size=batch_size, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True
+            )
+            query_embeddings = model.encode(
+                queries, batch_size=batch_size, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=True
+            )
         scores = model.similarity(query_embeddings, corpus_embeddings).to(device)
 
         # Keep only the range_max + max_positives highest scores. We offset by 1 to potentially include the positive pair