wanyaokun
70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
from typing import Any, List, Optional
|
|
from llama_index.core.postprocessor import SentenceTransformerRerank
|
|
from llama_index.core.schema import MetadataMode, NodeWithScore, QueryBundle
|
|
from llama_index.core.callbacks import CBEventType, EventPayload
|
|
from llama_index.core.bridge.pydantic import PrivateAttr
|
|
|
|
class OllamaRerank(SentenceTransformerRerank):
|
|
_score_threshold: float = PrivateAttr()
|
|
def __init__(
|
|
self,
|
|
top_n: int = 2,
|
|
model: str = "cross-encoder/stsb-distilroberta-base",
|
|
device: Optional[str] = None,
|
|
keep_retrieval_score: Optional[bool] = False,
|
|
score_threshold:float = 0.3
|
|
):
|
|
self._score_threshold = score_threshold
|
|
super().__init__(top_n,model,device,keep_retrieval_score)
|
|
|
|
@classmethod
|
|
def class_name(cls) -> str:
|
|
return "OllamaRerank"
|
|
|
|
def _postprocess_nodes(
|
|
self,
|
|
nodes: List[NodeWithScore],
|
|
query_bundle: Optional[QueryBundle] = None,
|
|
) -> List[NodeWithScore]:
|
|
if query_bundle is None:
|
|
raise ValueError("Missing query bundle in extra info.")
|
|
if len(nodes) == 0:
|
|
return []
|
|
|
|
query_and_nodes = [
|
|
(
|
|
query_bundle.query_str,
|
|
node.node.get_content(metadata_mode=MetadataMode.EMBED),
|
|
)
|
|
for node in nodes
|
|
]
|
|
|
|
with self.callback_manager.event(
|
|
CBEventType.RERANKING,
|
|
payload={
|
|
EventPayload.NODES: nodes,
|
|
EventPayload.MODEL_NAME: self.model,
|
|
EventPayload.QUERY_STR: query_bundle.query_str,
|
|
EventPayload.TOP_K: self.top_n,
|
|
},
|
|
) as event:
|
|
scores = self._model.predict(query_and_nodes)
|
|
|
|
assert len(scores) == len(nodes)
|
|
|
|
for node, score in zip(nodes, scores):
|
|
if self.keep_retrieval_score:
|
|
node.node.metadata["retrieval_score"] = node.score
|
|
node.score = score
|
|
|
|
for i in range(len(nodes)-1,-1,-1):
|
|
node = nodes[i]
|
|
if node.score < self._score_threshold:
|
|
nodes.remove(node)
|
|
|
|
new_nodes = sorted(nodes, key=lambda x: -x.score if x.score else 0)[
|
|
: self.top_n
|
|
]
|
|
event.on_end(payload={EventPayload.NODES: new_nodes})
|
|
|
|
return new_nodes |