Skip to content
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.

Sign up for GitHub

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

Adaptation for llama-index > 0.10 #87

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Adaptation for llama-index > 0.10 #87

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 10 additions & 9 deletions BCEmbedding/tools/llama_index/bce_rerank.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,20 +2,20 @@
@Description:
@Author: shenlei
@Date: 2024-01-15 14:15:30
@LastEditTime: 2024-03-04 16:00:52
@LastEditors: shenlei
@LastEditTime: 2024-09-06 13:25:04
@LastEditors: Innary
'''
from typing import Any, List, Optional

from pydantic.v1 import Field, PrivateAttr
from llama_index.callbacks import CBEventType, EventPayload
from llama_index.postprocessor.types import BaseNodePostprocessor
from llama_index.schema import MetadataMode, NodeWithScore, QueryBundle
from llama_index.utils import infer_torch_device
from llama_index.core.bridge.pydantic import Field, PrivateAttr
from llama_index.core.callbacks import CBEventType, EventPayload
from llama_index.core.postprocessor.types import BaseNodePostprocessor
from llama_index.core.schema import MetadataMode, NodeWithScore, QueryBundle
from llama_index.core.utils import infer_torch_device


class BCERerank(BaseNodePostprocessor):
model: str = Field(ddescription="Sentence transformer model name.")
model: str = Field(description="Sentence transformer model name.")
top_n: int = Field(description="Number of nodes to return sorted by score.")
_model: Any = PrivateAttr()

Expand All @@ -33,9 +33,10 @@ def __init__(
"Cannot import `BCEmbedding` package,",
"please `pip install BCEmbedding>=0.1.2`",
)
self._model = RerankerModel(model_name_or_path=model, device=device, **kwargs)
device = infer_torch_device() if device is None else device
super().__init__(top_n=top_n, model=model, device=device)
self._model = RerankerModel(model_name_or_path=model, device=device, **kwargs)


@classmethod
def class_name(cls) -> str:
Expand Down