diff --git a/docs/_src/api/api/ranker.md b/docs/_src/api/api/ranker.md index 6a84f1cdd..2b954800b 100644 --- a/docs/_src/api/api/ranker.md +++ b/docs/_src/api/api/ranker.md @@ -64,7 +64,9 @@ class SentenceTransformersRanker(BaseRanker) Sentence Transformer based pre-trained Cross-Encoder model for Document Re-ranking (https://huggingface.co/cross-encoder). Re-Ranking can be used on top of a retriever to boost the performance for document search. This is particularly useful if the retriever has a high recall but is bad in sorting the documents by relevance. -SentenceTransformerRanker handles Cross-Encoder models that use a single logit as similarity score. +SentenceTransformerRanker handles Cross-Encoder models + - use a single logit as similarity score e.g. cross-encoder/ms-marco-MiniLM-L-12-v2 + - use two output logits (no_answer, has_answer) e.g. deepset/gbert-base-germandpr-reranking https://www.sbert.net/docs/pretrained-models/ce-msmarco.html#usage-with-transformers | With a SentenceTransformersRanker, you can: diff --git a/haystack/modeling/model/optimization.py b/haystack/modeling/model/optimization.py index e84e49dbc..b8221fd95 100644 --- a/haystack/modeling/model/optimization.py +++ b/haystack/modeling/model/optimization.py @@ -303,7 +303,7 @@ def optimize_model(model, device, local_rank, optimizer=None, distributed=False, find_unused_parameters=True) elif torch.cuda.device_count() > 1 and device.type == "cuda": - model = WrappedDataParallel(model) + model = WrappedDataParallel(model) if not isinstance(model, DataParallel) else WrappedDataParallel(model.module) logger.info("Multi-GPU Training via DataParallel") return model, optimizer diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index 22673edd0..9f058cf7e 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -413,7 +413,8 @@ class DensePassageRetriever(BaseRetriever): self.query_tokenizer.save_pretrained(f"{save_dir}/{query_encoder_save_dir}") self.passage_tokenizer.save_pretrained(f"{save_dir}/{passage_encoder_save_dir}") - self.model = DataParallel(self.model, device_ids=self.devices) + if len(self.devices) > 1 and not isinstance(self.model, DataParallel): + self.model = DataParallel(self.model, device_ids=self.devices) def save(self, save_dir: Union[Path, str], query_encoder_dir: str = "query_encoder", passage_encoder_dir: str = "passage_encoder"):