Saltar al contenido principal

SentenceTransformerCrossEncoderRetriever

RetrieverModel
DashAI.back.models.RAG.retrievers.cross_encoder.SentenceTransformerCrossEncoderRetriever

Cross-encoder re-ranker powered by SentenceTransformer models.

Retrieves candidates from a single child retriever and re-ranks them using a SentenceTransformer CrossEncoder <https://sbert.net/docs/cross_encoder/pretrained_models.html>_ model that scores (query, chunk) pairs jointly.

The child retriever (the ranker) defines how many candidates are fetched via its own top_k; this cross-encoder only selects top_k of them after re-ranking.

Example usage::

Retrieve 15 candidates with BM25 (child's own top_k), keep top 5

after cross-encoding

config = { "model_name": "cross-encoder/ms-marco-MiniLM-L-6-v2", "top_k": 5, "children": [ { "component": "BM25Retriever", "params": {"...": "...", "top_k": 15}, } ], }

Parameters

model_name : string, default=cross-encoder/ms-marco-MiniLM-L-6-v2
Pre-trained SentenceTransformer cross-encoder model to use.
top_k : integer, default=5
Final number of chunks to return after re-ranking. The candidate set size is determined by the child retriever's own top_k.
children : array, default=[]
The child retriever whose candidates will be re-ranked.

Methods

get_metadata(cls) -> Dict[str, Any]

Defined on SentenceTransformerCrossEncoderRetriever

Return UI metadata including the declarative operation summary.

inject_infra(self, env_RAG_path: str, chunks: Dict[int, Dict[int, DashAI.back.models.RAG.documents.chunk.Chunk]], persistence: Any) -> None

Defined on SentenceTransformerCrossEncoderRetriever

Inject runtime infrastructure and load the cross-encoder model.

load(self, filename: str = '') -> None

Defined on SentenceTransformerCrossEncoderRetriever

Download and load the cross-encoder model from HuggingFace Hub.

add(self, child: DashAI.back.models.RAG.retrievers.retriever_model.RetrieverModel) -> None

Defined on CompositeRetriever

Add a child retriever.

calculate_metrics(self, split: DashAI.back.core.enums.metrics.SplitEnum = <SplitEnum.VALIDATION: 'validation'>, level: DashAI.back.core.enums.metrics.LevelEnum = <LevelEnum.LAST: 'last'>, log_index: int = None, x_data: 'DashAIDataset' = None, y_data: 'DashAIDataset' = None, fold_index: int = None, inner_fold_index: int = None)

Defined on BaseModel

Calculate and save metrics for a given data split and level.

Parameters

split : SplitEnum
The data split to evaluate (TRAIN, VALIDATION, or TEST). Defaults to SplitEnum.VALIDATION.
level : LevelEnum
The metric granularity level (LAST, TRIAL, STEP, or BATCH). Defaults to LevelEnum.LAST.
log_index : int, optional
Explicit step index for the metric entry. If None, the next step index is computed automatically. Defaults to None.
x_data : DashAIDataset, optional
Input features. If None, the dataset stored in the model for the given split is used. Defaults to None.
y_data : DashAIDataset, optional
Target labels. If None, the labels stored in the model for the given split are used. Defaults to None.

compute_metrics(self, split: DashAI.back.core.enums.metrics.SplitEnum = <SplitEnum.TEST: 'test'>, x_data: 'DashAIDataset' = None, y_data: 'DashAIDataset' = None) -> Dict[str, float]

Defined on BaseModel

Calculate and return metric scores for a given data split.

Parameters

split : SplitEnum
The data split to evaluate (TRAIN, VALIDATION, or TEST). Defaults to SplitEnum.VALIDATION.
x_data : DashAIDataset, optional
Input features. If None, the dataset stored in the model for the given split is used. Defaults to None.
y_data : DashAIDataset, optional
Target labels. If None, the labels stored in the model for the given split are used. Defaults to None.

Returns

Dict[str, float]
A dictionary mapping metric names to their computed scores.

get_children(self) -> List[DashAI.back.models.RAG.retrievers.retriever_model.RetrieverModel]

Defined on CompositeRetriever

Return a copy of the children list.

get_chunk_vectors(self, chunk_ids: List[int]) -> numpy.ndarray

Defined on CrossEncoderRetriever

Return chunk vectors from the child retriever.

get_credential(self, name: str)

Defined on ConfigObject

Resolve a registered credential component by name.

Parameters

name : str
Credential component class name (e.g. "HuggingFaceCredential").

Returns

BaseCredential
An instance of the requested credential component.

get_id(self) -> int | None

Defined on RetrieverModel

Return the database ID of this retriever, or None.

get_schema(cls) -> dict

Defined on ConfigObject

Generates the component related Json Schema.

Returns

dict
Dictionary representing the Json Schema of the component.

init_model(self) -> None

Defined on CrossEncoderRetriever

Initialize the model by loading it.

predict_prepared(self, features: Any) -> Any

Defined on BaseModel

Predict from data that is already in this model's feature space.

Parameters

features : pandas.DataFrame or numpy.ndarray
Feature matrix as returned by prepare_dataset(..., is_fit=False). No further preparation is applied to it.

Returns

Any
The same kind of output as predict: predicted values for regressors, class probabilities for DashAI classifiers.

predict_proba_prepared(self, features: Any) -> Any

Defined on BaseModel

Return class probabilities for data already in the feature space.

Parameters

features : pandas.DataFrame or numpy.ndarray
Feature matrix as returned by prepare_dataset(..., is_fit=False).

Returns

numpy.ndarray
Array of shape (n_samples, n_classes) with class probabilities.

prepare_dataset(self, dataset: 'DashAIDataset', is_fit: bool = False) -> 'DashAIDataset'

Defined on BaseModel

Hook for model specific preprocessing of input features.

Parameters

dataset : DashAIDataset
The input dataset to preprocess.
is_fit : bool
Whether the call is part of a fitting phase. Defaults to False.

Returns

DashAIDataset
The preprocessed dataset ready to be fed into the model.

prepare_output(self, dataset: 'DashAIDataset', is_fit: bool = False) -> 'DashAIDataset'

Defined on BaseModel

Hook for model-specific preprocessing of output targets.

Parameters

dataset : DashAIDataset
The output dataset (target labels) to preprocess.
is_fit : bool
Whether the call is part of a fitting phase. Defaults to False.

Returns

DashAIDataset
The preprocessed output dataset.

remove(self, child: DashAI.back.models.RAG.retrievers.retriever_model.RetrieverModel) -> None

Defined on CompositeRetriever

Remove a child retriever.

retrieve(self, query: str, top_k: int | None = None, **kwargs) -> List[DashAI.back.models.RAG.documents.chunk.Chunk]

Defined on CrossEncoderRetriever

Retrieve and re-rank chunks using the cross-encoder.

save(self, filename: str = '') -> None

Defined on CrossEncoderRetriever

Persist the cross-encoder retriever's state.

score_chunks(self, chunk_ids: List[int], query: str) -> List[Tuple[int, float]]

Defined on CrossEncoderRetriever

Score a set of chunk IDs against the query with the cross-encoder.

set_id(self, id: int) -> None

Defined on RetrieverModel

Assign a database ID to this retriever.

train(self, **kwargs)

Defined on RetrieverModel

Train the retriever on the injected chunks.

validate_and_transform(self, raw_data: dict) -> dict

Defined on ConfigObject

It takes the data given by the user to initialize the model and returns it with all the objects that the model needs to work.

Parameters

raw_data : dict
A dictionary with the data provided by the user to initialize the model.

Returns

dict
A validated dictionary with the necessary objects.