-
Notifications
You must be signed in to change notification settings - Fork 1
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat(document-search): LiteLLM Reranker #109
base: main
Are you sure you want to change the base?
feat(document-search): LiteLLM Reranker #109
Conversation
Code Coverage Summary
Diff against main
Results for commit: 5efc234 Minimum allowed coverage is ♻️ This comment has been updated with latest results |
Trivy scanning results. .venv/lib/python3.10/site-packages/PyJWT-2.9.0.dist-info/METADATA (secrets)Total: 1 (MEDIUM: 1, HIGH: 0, CRITICAL: 0) MEDIUM: JWT (jwt-token) .venv/lib/python3.10/site-packages/litellm/llms/huggingface_llms_metadata/hf_text_generation_models.txt (secrets)Total: 1 (MEDIUM: 0, HIGH: 0, CRITICAL: 1) CRITICAL: HuggingFace (hugging-face-access-token) .venv/lib/python3.10/site-packages/litellm/proxy/_types.py (secrets)Total: 1 (MEDIUM: 1, HIGH: 0, CRITICAL: 0) MEDIUM: Slack (slack-web-hook) |
from ragbits.document_search.documents.element import Element | ||
|
||
|
||
class Reranker(abc.ABC): | ||
class Reranker(BaseModel, abc.ABC): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
maybe a stupid question, but why is pydantic
needed here?
|
||
documents = [chunk.content if isinstance(chunk, TextElement) else None for chunk in chunks] | ||
|
||
response = litellm.rerank( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
maybe it would be better to use arerank
here: https://docs.litellm.ai/docs/rerank#async-usage, what do you think?
|
||
class LiteLLMReranker(Reranker): | ||
""" | ||
A LiteLLM reranker for providers such as Cohere, Together AI, Azure AI. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think it might be helpful to describe a bit these attributes below or link the documentation, I couldn't find litellm
, but cohere
is probably the same https://docs.cohere.com/v2/reference/rerank#request
from ragbits.document_search.retrieval.rerankers.base import Reranker | ||
|
||
|
||
class LiteLLMReranker(Reranker): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
i think you should add it here: https://github.com/deepsense-ai/ragbits/blob/main/packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/__init__.py#L9 for it to work with config like this:
"reranker": {
"type": "LiteLLMReranker",
"config": {"model": "cohere/rerank-english-v3.0"},
}
No description provided.