mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-07 14:03:26 +00:00
langchain[minor], community[minor]: add CrossEncoderReranker with HuggingFaceCrossEncoder and SagemakerEndpointCrossEncoder (#13687)
- **Description:** Support reranking based on cross encoder models available from HuggingFace. - Added `CrossEncoder` schema - Implemented `HuggingFaceCrossEncoder` and `SagemakerEndpointCrossEncoder` - Implemented `CrossEncoderReranker` that performs similar functionality to `CohereRerank` - Added `cross-encoder-reranker.ipynb` to demonstrate how to use it. Please let me know if anything else needs to be done to make it visible on the table-of-contents navigation bar on the left, or on the card list on [retrievers documentation page](https://python.langchain.com/docs/integrations/retrievers). - **Issue:** N/A - **Dependencies:** None other than the existing ones. --------- Co-authored-by: Kenny Choe <kchoe@amazon.com> Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
@@ -0,0 +1,63 @@
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
from langchain_core.pydantic_v1 import BaseModel, Extra, Field
|
||||
|
||||
from langchain_community.cross_encoders.base import BaseCrossEncoder
|
||||
|
||||
DEFAULT_MODEL_NAME = "BAAI/bge-reranker-base"
|
||||
|
||||
|
||||
class HuggingFaceCrossEncoder(BaseModel, BaseCrossEncoder):
|
||||
"""HuggingFace cross encoder models.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
|
||||
|
||||
model_name = "BAAI/bge-reranker-base"
|
||||
model_kwargs = {'device': 'cpu'}
|
||||
hf = HuggingFaceCrossEncoder(
|
||||
model_name=model_name,
|
||||
model_kwargs=model_kwargs
|
||||
)
|
||||
"""
|
||||
|
||||
client: Any #: :meta private:
|
||||
model_name: str = DEFAULT_MODEL_NAME
|
||||
"""Model name to use."""
|
||||
model_kwargs: Dict[str, Any] = Field(default_factory=dict)
|
||||
"""Keyword arguments to pass to the model."""
|
||||
|
||||
def __init__(self, **kwargs: Any):
|
||||
"""Initialize the sentence_transformer."""
|
||||
super().__init__(**kwargs)
|
||||
try:
|
||||
import sentence_transformers
|
||||
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Could not import sentence_transformers python package. "
|
||||
"Please install it with `pip install sentence-transformers`."
|
||||
) from exc
|
||||
|
||||
self.client = sentence_transformers.CrossEncoder(
|
||||
self.model_name, **self.model_kwargs
|
||||
)
|
||||
|
||||
class Config:
|
||||
"""Configuration for this pydantic object."""
|
||||
|
||||
extra = Extra.forbid
|
||||
|
||||
def score(self, text_pairs: List[Tuple[str, str]]) -> List[float]:
|
||||
"""Compute similarity scores using a HuggingFace transformer model.
|
||||
|
||||
Args:
|
||||
text_pairs: The list of text text_pairs to score the similarity.
|
||||
|
||||
Returns:
|
||||
List of scores, one for each pair.
|
||||
"""
|
||||
scores = self.client.predict(text_pairs)
|
||||
return scores
|
Reference in New Issue
Block a user