Custom rerankers support #496
Conversation
📝 WalkthroughWalkthroughThe changes introduce a new class method Possibly related PRs
Suggested reviewers
Tip ⚡🧪 Multi-step agentic review comment chat (experimental)
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
fastembed/rerank/cross_encoder/custom_reranker_model.py (3)
10-12: Class-level SUPPORTED_MODELS could lead to issues.The
SUPPORTED_MODELSlist is defined at the class level and can be shared across all instances of the class. Consider alternatives:
- If this is intentionally a class-level shared registry, add a comment explaining this design.
- Otherwise, consider initializing in
__init__if model lists should be instance-specific.Also, missing docstring for the class that explains its purpose and usage.
class CustomCrossEncoderModel(OnnxTextCrossEncoder): + """ + Custom cross-encoder model that allows adding and managing custom reranker models. + + This class extends OnnxTextCrossEncoder to provide functionality for registering + and using custom reranker models via a class-level registry of model descriptions. + """ SUPPORTED_MODELS: list[DenseModelDescription] = [] + # Class-level registry intentionally shared across all instances
13-38: Constructor implementation is appropriate.The constructor properly forwards all parameters to the parent class, maintaining the expected interface.
Consider adding a docstring to document the parameters and their usage.
def __init__( self, model_name: str, cache_dir: Optional[str] = None, threads: Optional[int] = None, providers: Optional[Sequence[OnnxProvider]] = None, cuda: bool = False, device_ids: Optional[list[int]] = None, lazy_load: bool = False, device_id: Optional[int] = None, specific_model_path: Optional[str] = None, **kwargs: Any, ): + """ + Initialize a custom cross-encoder model. + + Args: + model_name: Name of the model to use + cache_dir: Directory to cache models + threads: Number of threads to use for inference + providers: ONNX execution providers + cuda: Whether to use CUDA + device_ids: List of device IDs to use + lazy_load: Whether to load the model lazily + device_id: Specific device ID to use + specific_model_path: Path to a specific model file + **kwargs: Additional arguments passed to the parent class + """ super().__init__( model_name=model_name, cache_dir=cache_dir, threads=threads, providers=providers, cuda=cuda, device_ids=device_ids, lazy_load=lazy_load, device_id=device_id, specific_model_path=specific_model_path, **kwargs, )
43-49: Add docstring and validation to add_model method.The
add_modelmethod lacks a docstring and input validation. Consider adding:
- A descriptive docstring that explains the method's purpose and parameters
- Validation to ensure the provided model description is valid
@classmethod def add_model( cls, model_description: DenseModelDescription, ) -> None: + """ + Add a custom model to the list of supported models. + + Args: + model_description: Description of the model to add + + Raises: + ValueError: If the model description is invalid or a model with the same name already exists + """ + # Validate model description + if not model_description.source or not model_description.model_file: + raise ValueError("Model description must include source and model file") + + # Check for duplicate model names + if any(model.name == model_description.name for model in cls.SUPPORTED_MODELS): + raise ValueError(f"Model with name '{model_description.name}' already exists") + cls.SUPPORTED_MODELS.append(model_description)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
fastembed/rerank/cross_encoder/__init__.py(1 hunks)fastembed/rerank/cross_encoder/custom_reranker_model.py(1 hunks)
🧰 Additional context used
🪛 GitHub Actions: type-checkers
fastembed/rerank/cross_encoder/custom_reranker_model.py
[error] 40-40: Return type 'list[DenseModelDescription]' of '_list_supported_models' incompatible with return type 'list[BaseModelDescription]' in supertype 'OnnxTextCrossEncoder' [override]
[error] 40-40: Return type 'list[DenseModelDescription]' of '_list_supported_models' incompatible with return type 'list[BaseModelDescription]' in supertype 'ModelManagement' [override]
⏰ Context from checks skipped due to timeout of 90000ms (16)
- GitHub Check: Python 3.13.x on windows-latest test
- GitHub Check: Python 3.13.x on macos-latest test
- GitHub Check: Python 3.13.x on ubuntu-latest test
- GitHub Check: Python 3.12.x on windows-latest test
- GitHub Check: Python 3.12.x on macos-latest test
- GitHub Check: Python 3.12.x on ubuntu-latest test
- GitHub Check: Python 3.11.x on windows-latest test
- GitHub Check: Python 3.11.x on macos-latest test
- GitHub Check: Python 3.11.x on ubuntu-latest test
- GitHub Check: Python 3.10.x on windows-latest test
- GitHub Check: Python 3.10.x on macos-latest test
- GitHub Check: Python 3.10.x on ubuntu-latest test
- GitHub Check: Python 3.9.x on windows-latest test
- GitHub Check: Python 3.9.x on macos-latest test
- GitHub Check: Python 3.9.x on ubuntu-latest test
- GitHub Check: Analyze (python)
🔇 Additional comments (4)
fastembed/rerank/cross_encoder/__init__.py (2)
2-2: Import statement looks good.The import of
CustomCrossEncoderModelfromfastembed.rerank.cross_encoder.custom_reranker_modelis correctly added.
4-4: Public interface update is appropriate.The
__all__list is properly updated to include the newCustomCrossEncoderModelclass, making it part of the module's public API.fastembed/rerank/cross_encoder/custom_reranker_model.py (2)
1-8: Imports are appropriate and organized well.The necessary imports for the implementation are included and well-organized, with standard library imports first, followed by package-specific imports.
39-42:Details
✅ Verification successful
Type error in _list_supported_models return annotation.
The pipeline failures indicate a type compatibility issue. The method overrides a parent method that returns
list[BaseModelDescription], but this implementation specifieslist[DenseModelDescription].@classmethod - def _list_supported_models(cls) -> list[DenseModelDescription]: + def _list_supported_models(cls) -> list["BaseModelDescription"]: return cls.SUPPORTED_MODELSThis assumes that
DenseModelDescriptionis a subclass ofBaseModelDescription. If not, you'll need to ensure type compatibility by other means.
🏁 Script executed:
#!/bin/bash # Check the inheritance hierarchy of DenseModelDescription grep -r "class DenseModelDescription" --include="*.py" . grep -r "BaseModelDescription" --include="*.py" .Length of output: 2686
Type Annotation Correction in
_list_supported_modelsThe issue is confirmed: the parent method expects a return type of
list[BaseModelDescription], while this override uses the more specificlist[DenseModelDescription]. SinceDenseModelDescriptionis indeed a subclass ofBaseModelDescription(as verified infastembed/common/model_description.py), the method signature should be updated to the parent’s annotation to ensure type compatibility.
- File:
fastembed/rerank/cross_encoder/custom_reranker_model.py- Lines: 39–42
- Change: Update the return annotation from
list[DenseModelDescription]tolist["BaseModelDescription"](using a string literal for forward reference) to correctly reflect the parent type.@classmethod - def _list_supported_models(cls) -> list[DenseModelDescription]: + def _list_supported_models(cls) -> list["BaseModelDescription"]: return cls.SUPPORTED_MODELS🧰 Tools
🪛 GitHub Actions: type-checkers
[error] 40-40: Return type 'list[DenseModelDescription]' of '_list_supported_models' incompatible with return type 'list[BaseModelDescription]' in supertype 'OnnxTextCrossEncoder' [override]
[error] 40-40: Return type 'list[DenseModelDescription]' of '_list_supported_models' incompatible with return type 'list[BaseModelDescription]' in supertype 'ModelManagement' [override]
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
fastembed/rerank/cross_encoder/text_cross_encoder.py (1)
135-167: Well-implemented custom model registration methodThe
add_custom_modelmethod is well-structured with appropriate parameter validation and error handling. It checks for existing models before registration and provides sensible defaults.Two minor suggestions:
- Consider adding docstring documentation for this public method, similar to other methods in the class
- The error message on line 151-152 refers to "CrossEncoderModel" but should probably say "TextCrossEncoder" for consistency
@classmethod def add_custom_model( cls, model: str, sources: ModelSource, dim: int, model_file: str = "onnx/model.onnx", description: str = "", license: str = "", size_in_gb: float = 0.0, additional_files: Optional[list[str]] = None, + ) -> None: + """Registers a custom cross-encoder model. + + Args: + model: Unique name for the custom model + sources: Model source information + dim: Model output dimension + model_file: Path to the model file, defaults to "onnx/model.onnx" + description: Optional model description + license: Optional license information + size_in_gb: Size of the model in GB + additional_files: List of additional files required by the model + + Raises: + ValueError: If a model with the given name is already registered + """ registered_models = cls._list_supported_models() for registered_model in registered_models: if model == registered_model.model: raise ValueError( - f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, " + f"Model {model} is already registered in TextCrossEncoder, if you still want to add this model, " f"please use another model name" )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
fastembed/rerank/cross_encoder/custom_reranker_model.py(1 hunks)fastembed/rerank/cross_encoder/text_cross_encoder.py(2 hunks)tests/test_custom_models.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- fastembed/rerank/cross_encoder/custom_reranker_model.py
⏰ Context from checks skipped due to timeout of 90000ms (15)
- GitHub Check: Python 3.13.x on windows-latest test
- GitHub Check: Python 3.13.x on macos-latest test
- GitHub Check: Python 3.13.x on ubuntu-latest test
- GitHub Check: Python 3.12.x on windows-latest test
- GitHub Check: Python 3.12.x on macos-latest test
- GitHub Check: Python 3.12.x on ubuntu-latest test
- GitHub Check: Python 3.11.x on windows-latest test
- GitHub Check: Python 3.11.x on macos-latest test
- GitHub Check: Python 3.11.x on ubuntu-latest test
- GitHub Check: Python 3.10.x on windows-latest test
- GitHub Check: Python 3.10.x on macos-latest test
- GitHub Check: Python 3.10.x on ubuntu-latest test
- GitHub Check: Python 3.9.x on windows-latest test
- GitHub Check: Python 3.9.x on macos-latest test
- GitHub Check: Python 3.9.x on ubuntu-latest test
🔇 Additional comments (4)
fastembed/rerank/cross_encoder/text_cross_encoder.py (3)
6-7: Appropriate addition of the custom reranker model importThe import for
CustomCrossEncoderModelis necessary for the registry update and the new functionality added to theTextCrossEncoderclass.
9-13: Clean expansion of model description importsAll necessary model description classes are now properly imported. The additional imports are required for the new
add_custom_modelmethod.
19-19: Good addition of CustomCrossEncoderModel to the registryAdding the model to the registry enables the TextCrossEncoder to recognize and use custom reranker models.
tests/test_custom_models.py (1)
10-11: Proper imports for testing custom rerankersThe imports for
CustomCrossEncoderModelandTextCrossEncoderare necessary for the new test function.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
fastembed/rerank/cross_encoder/custom_reranker_model.py (1)
41-47: Consider adding duplicate model validation.The current implementation allows adding duplicate models with the same name to the SUPPORTED_MODELS list. Consider adding validation to prevent duplicates.
@classmethod def add_model( cls, model_description: BaseModelDescription, ) -> None: + # Check if model with same name already exists + for existing_model in cls.SUPPORTED_MODELS: + if existing_model.model_name == model_description.model_name: + raise ValueError(f"Model with name '{model_description.model_name}' already exists") cls.SUPPORTED_MODELS.append(model_description)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
fastembed/rerank/cross_encoder/custom_reranker_model.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (16)
- GitHub Check: Python 3.13.x on windows-latest test
- GitHub Check: Python 3.13.x on macos-latest test
- GitHub Check: Python 3.13.x on ubuntu-latest test
- GitHub Check: Python 3.12.x on windows-latest test
- GitHub Check: Python 3.12.x on macos-latest test
- GitHub Check: Python 3.12.x on ubuntu-latest test
- GitHub Check: Python 3.11.x on windows-latest test
- GitHub Check: Python 3.11.x on macos-latest test
- GitHub Check: Python 3.11.x on ubuntu-latest test
- GitHub Check: Python 3.10.x on windows-latest test
- GitHub Check: Python 3.10.x on macos-latest test
- GitHub Check: Python 3.10.x on ubuntu-latest test
- GitHub Check: Python 3.9.x on windows-latest test
- GitHub Check: Python 3.9.x on macos-latest test
- GitHub Check: Python 3.13 test
- GitHub Check: Python 3.9.x on ubuntu-latest test
🔇 Additional comments (4)
fastembed/rerank/cross_encoder/custom_reranker_model.py (4)
8-9: Nice implementation of a custom cross encoder model registry.The class design allows users to dynamically register their own cross encoder models, effectively extending the library's capabilities without modifying core code.
11-35: Constructor implementation looks good.The constructor properly inherits from OnnxTextCrossEncoder and passes all the necessary parameters to the parent class. The type hints are well-defined, and the method follows best practices by accepting **kwargs to handle any additional parameters.
37-40: Simple and effective implementation.This class method provides a clean way to access the list of supported models.
9-9:Details
❓ Verification inconclusive
Consider thread-safety for the model registry.
The
SUPPORTED_MODELSclass variable is mutable and could potentially be accessed from multiple threads. If this code will be used in a multi-threaded environment, consider using thread-safe data structures or implementing synchronization.
🏁 Script executed:
#!/bin/bash # Check if the code might be used in a multi-threaded context rg -A 5 -B 5 "Thread|thread|concurrent|Concurrent|Lock|lock|Semaphore|semaphore" --glob "**/*.py"Length of output: 77618
Attention: Thread-Safety for Model Registry
In
fastembed/rerank/cross_encoder/custom_reranker_model.py, the class variableSUPPORTED_MODELS: list[BaseModelDescription] = []is defined as a mutable list without any synchronization. While it appears that this registry is populated during initialization (and not modified thereafter), the broader codebase shows several components that operate in a multi-threaded context. To avoid potential race conditions if the registry is ever updated concurrently, please verify the following:
- If modifications are expected at runtime: Consider using synchronization mechanisms (e.g., a threading lock) or a thread-safe data structure to guard updates.
- If the registry is static post-initialization: Converting it to an immutable tuple after population might be a safer alternative.
Please review the intended usage of
SUPPORTED_MODELSunder concurrent scenarios and adjust its implementation accordingly.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
fastembed/rerank/cross_encoder/text_cross_encoder.py (1)
133-164: Minor naming enhancement suggested.The
add_custom_modelmethod is well-structured and clearly checks for duplicates before adding a new model. However, consider renaming the parametermodeltomodel_namefor readability, since “model” might be confused with an instantiated model object.- def add_custom_model( - cls, - model: str, + def add_custom_model( + cls, + model_name: str, sources: ModelSource, ... ): ... for registered_model in registered_models: - if model == registered_model.model: + if model_name == registered_model.model: ... ... - CustomCrossEncoderModel.add_model( + CustomCrossEncoderModel.add_model( BaseModelDescription( - model=model, + model=model_name, ... ) )tests/test_custom_models.py (1)
77-114: Rename to clarify meaning.The test logic is solid. However, variable
embeddingsmight be misleading because the cross-encoder returns a scalar score for each pair. Consider renaming it toscoresoroutputsfor clarity.- embeddings = np.stack(scores, axis=0) - assert embeddings.shape == (2,) - assert np.allclose(embeddings, canonical_vector, atol=1e-3) + scored_outputs = np.stack(scores, axis=0) + assert scored_outputs.shape == (2,) + assert np.allclose(scored_outputs, canonical_vector, atol=1e-3)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
fastembed/rerank/cross_encoder/text_cross_encoder.py(2 hunks)tests/test_custom_models.py(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (16)
- GitHub Check: Python 3.13.x on windows-latest test
- GitHub Check: Python 3.13.x on macos-latest test
- GitHub Check: Python 3.13.x on ubuntu-latest test
- GitHub Check: Python 3.12.x on windows-latest test
- GitHub Check: Python 3.12.x on macos-latest test
- GitHub Check: Python 3.12.x on ubuntu-latest test
- GitHub Check: Python 3.11.x on windows-latest test
- GitHub Check: Python 3.11.x on macos-latest test
- GitHub Check: Python 3.11.x on ubuntu-latest test
- GitHub Check: Python 3.10.x on windows-latest test
- GitHub Check: Python 3.10.x on macos-latest test
- GitHub Check: Python 3.13 test
- GitHub Check: Python 3.10.x on ubuntu-latest test
- GitHub Check: Python 3.9.x on windows-latest test
- GitHub Check: Python 3.9.x on macos-latest test
- GitHub Check: Python 3.9.x on ubuntu-latest test
🔇 Additional comments (7)
fastembed/rerank/cross_encoder/text_cross_encoder.py (3)
6-6: Module import looks good.Including
CustomCrossEncoderModelas part of the imports is correct and aligns with the new functionality for custom rerankers.
9-12: Imports are valid.The added imports for
ModelSourceandBaseModelDescriptionare necessary for the new custom model registration logic.
18-18: Registration extension is appropriate.Adding
CustomCrossEncoderModelto theCROSS_ENCODER_REGISTRYensures theTextCrossEncodercan instantiate custom cross-encoder models.tests/test_custom_models.py (4)
6-11: New imports are valid.Importing
PoolingType,ModelSource,DenseModelDescription, andBaseModelDescriptionis essential for describing and registering custom models.
15-16: Importing cross-encoder classes is correct.Bringing in
CustomCrossEncoderModelandTextCrossEncoderaligns with the tests for custom cross-encoder functionality.
24-27: Fixture cleanup is consistent.Resetting
CustomCrossEncoderModel.SUPPORTED_MODELSbefore and after tests ensures a clean environment across tests.
209-232: Validation of duplicate model addition is correct.The test correctly ensures a
ValueErroris raised when attempting to add an already registered cross-encoder.
|
This would be nice |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
fastembed/rerank/cross_encoder/text_cross_encoder.py (1)
134-163: Well-implemented custom model registration methodThe
add_custom_modelmethod is well-structured with:
- Proper validation to prevent duplicate model registration
- Clear error messaging
- Appropriate parameter forwarding to
CustomTextCrossEncoder.add_modelHowever, there's a minor inconsistency in the error message.
- f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, " + f"Model {model} is already registered in TextCrossEncoder, if you still want to add this model, "
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
fastembed/rerank/cross_encoder/custom_text_cross_encoder.py(1 hunks)fastembed/rerank/cross_encoder/text_cross_encoder.py(2 hunks)tests/test_custom_models.py(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (16)
- GitHub Check: Python 3.13.x on windows-latest test
- GitHub Check: Python 3.13.x on macos-latest test
- GitHub Check: Python 3.13.x on ubuntu-latest test
- GitHub Check: Python 3.12.x on windows-latest test
- GitHub Check: Python 3.12.x on macos-latest test
- GitHub Check: Python 3.12.x on ubuntu-latest test
- GitHub Check: Python 3.11.x on windows-latest test
- GitHub Check: Python 3.11.x on macos-latest test
- GitHub Check: Python 3.11.x on ubuntu-latest test
- GitHub Check: Python 3.10.x on windows-latest test
- GitHub Check: Python 3.10.x on macos-latest test
- GitHub Check: Python 3.10.x on ubuntu-latest test
- GitHub Check: Python 3.9.x on windows-latest test
- GitHub Check: Python 3.9.x on macos-latest test
- GitHub Check: Python 3.9.x on ubuntu-latest test
- GitHub Check: Python 3.13 test
🔇 Additional comments (10)
fastembed/rerank/cross_encoder/custom_text_cross_encoder.py (4)
8-9: Good implementation of model registryThe class-level attribute
SUPPORTED_MODELSis appropriately initialized as an empty list to store custom model descriptions.
11-35: Implementation follows good inheritance patternsThe constructor correctly passes all parameters to the parent class, maintaining consistent initialization behavior with the base
OnnxTextCrossEncoderclass.
37-40: Well-implemented model listing methodThis method correctly returns the class-level
SUPPORTED_MODELSlist, allowing the registry to be queried.
41-47: Clean model registration implementationThe
add_modelclass method appropriately appends new model descriptions to the registry, enabling dynamic model registration.fastembed/rerank/cross_encoder/text_cross_encoder.py (2)
6-12: Appropriate imports for new functionalityThe necessary imports for
CustomTextCrossEncoderand model description classes are correctly added.
16-19: Registry properly updatedThe
CustomTextCrossEncoderis correctly added to theCROSS_ENCODER_REGISTRY, enabling discovery of custom models.tests/test_custom_models.py (4)
6-16: Imports are properly organizedAll necessary imports for the new cross-encoder functionality are correctly added, including model description classes and the custom cross-encoder implementation.
22-27: Fixture correctly updated to reset custom modelsThe
restore_custom_models_fixturehas been appropriately updated to resetCustomTextCrossEncoder.SUPPORTED_MODELSboth before and after tests, ensuring test isolation.
77-114: Comprehensive test for custom cross-encoder modelThe test function is well-implemented and covers all essential aspects:
- Model registration using
add_custom_model- Verification of model properties
- Creating a model instance and using it to rerank pairs
- Checking the output shape and values
- Cleanup in CI environments
I see you've included
additional_files=["onnx/model.onnx_data"]as recommended in previous comments.
209-232: Robust duplicate detection testThis test properly verifies that:
- Adding an existing base model raises an appropriate error
- A new custom model can be added successfully
- Attempting to add the same custom model again raises an error
This ensures the validation logic in
add_custom_modelworks correctly.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tests/test_custom_models.py (2)
77-112: Implementation for custom cross-encoder testing looks goodThe test follows a similar pattern to the existing custom text embedding test, which ensures consistency in the codebase. It properly:
- Sets up test data with a custom model
- Registers the model using the new API
- Verifies the model is correctly registered
- Creates a model instance and tests functionality
- Cleans up resources in CI environments
However, I notice this test uses "Xenova/ms-marco-MiniLM-L-4-v2" while the PR objectives mentioned "viplao5/bge-reranker-v2-m3-onnx". This is likely a reasonable testing choice for faster/simpler tests, but you might consider adding another test with the actual model mentioned in the PR objectives for complete validation.
84-89: Consider adding support for additional_files parameterPrevious review comments indicated that some ONNX models (like "viplao5/bge-reranker-v2-m3-onnx") require additional files such as "model.onnx_data". While the current test model may not need this, it would be good to ensure the
add_custom_modelimplementation supports this parameter and that it's tested.TextCrossEncoder.add_custom_model( custom_model_name, model_file="onnx/model.onnx", sources=source, size_in_gb=size_in_gb, + additional_files=[], # Add support for models that require additional files )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/test_custom_models.py(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (16)
- GitHub Check: Python 3.13.x on windows-latest test
- GitHub Check: Python 3.13.x on macos-latest test
- GitHub Check: Python 3.13.x on ubuntu-latest test
- GitHub Check: Python 3.12.x on windows-latest test
- GitHub Check: Python 3.12.x on macos-latest test
- GitHub Check: Python 3.12.x on ubuntu-latest test
- GitHub Check: Python 3.11.x on windows-latest test
- GitHub Check: Python 3.11.x on macos-latest test
- GitHub Check: Python 3.11.x on ubuntu-latest test
- GitHub Check: Python 3.10.x on windows-latest test
- GitHub Check: Python 3.10.x on macos-latest test
- GitHub Check: Python 3.10.x on ubuntu-latest test
- GitHub Check: Python 3.9.x on windows-latest test
- GitHub Check: Python 3.9.x on macos-latest test
- GitHub Check: Python 3.9.x on ubuntu-latest test
- GitHub Check: Python 3.13 test
🔇 Additional comments (3)
tests/test_custom_models.py (3)
15-16: Good addition of necessary importsThese new imports correctly include the custom cross-encoder classes needed for testing the new reranker functionality.
24-27: Good update of test fixtureThe fixture has been properly modified to reset
CustomTextCrossEncoder.SUPPORTED_MODELSbefore and after each test, ensuring test isolation and preventing test interference.
207-229: Good error handling testThis test properly verifies that attempting to register an already existing cross-encoder model raises a
ValueError. The test has good coverage of both built-in models and custom models that have been added. This ensures the API behaves correctly and provides appropriate error messages.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
README.md (2)
193-194: Improve Clarity in Custom Model Extension Explanation
The sentence is clear, but consider rewording it for even better readability. For example, you could change it to:-Text cross encoders can also be extended with models which are not in the list of supported models. +You can extend text cross encoders to support custom models not originally included in our list of supported models.This revision makes it more direct and easier for users to understand the extensibility of custom models.
195-208: Demonstrative Usage Example for Custom Models
The code snippet clearly demonstrates how to register and utilize a custom model using theadd_custom_modelmethod and thererank_pairsAPI. It effectively shows the complete flow—from importing the required classes to instantiating the model and retrieving rerank scores.For additional clarity, consider adding inline comments that briefly explain each step (e.g., “Register custom model”, “Instantiate model”, “Rerank pairs”) to guide new users through the process.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
README.md(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (16)
- GitHub Check: Python 3.13.x on windows-latest test
- GitHub Check: Python 3.13.x on macos-latest test
- GitHub Check: Python 3.13.x on ubuntu-latest test
- GitHub Check: Python 3.12.x on windows-latest test
- GitHub Check: Python 3.12.x on macos-latest test
- GitHub Check: Python 3.12.x on ubuntu-latest test
- GitHub Check: Python 3.11.x on windows-latest test
- GitHub Check: Python 3.11.x on macos-latest test
- GitHub Check: Python 3.11.x on ubuntu-latest test
- GitHub Check: Python 3.10.x on windows-latest test
- GitHub Check: Python 3.10.x on macos-latest test
- GitHub Check: Python 3.10.x on ubuntu-latest test
- GitHub Check: Python 3.9.x on windows-latest test
- GitHub Check: Python 3.9.x on macos-latest test
- GitHub Check: Python 3.9.x on ubuntu-latest test
- GitHub Check: Python 3.13 test
Added support of custom rerankers. It unlocks 494 functionality as: