From d3a2aec2cf549352e5e76ab8425797729f1a2471 Mon Sep 17 00:00:00 2001 From: "xianzhe.xxz" Date: Fri, 26 Jul 2024 17:16:49 +0800 Subject: [PATCH] fix dashscope rerank:add alpha support, when alpha=1.0, only use vector similarity, when alpha=0.0, only use bm25, when alpha=None, use rrf to fuse the results --- memoryscope/models/llama_index_rank_model.py | 7 +++-- .../storage/llama_index_es_memory_store.py | 5 +++- .../storage/llama_index_sync_elasticsearch.py | 26 ++++++++++++++++--- 3 files changed, 31 insertions(+), 7 deletions(-) diff --git a/memoryscope/models/llama_index_rank_model.py b/memoryscope/models/llama_index_rank_model.py index 9e74cdbd..38debfd7 100644 --- a/memoryscope/models/llama_index_rank_model.py +++ b/memoryscope/models/llama_index_rank_model.py @@ -35,12 +35,13 @@ class LlamaIndexRankModel(BaseModel): documents = [documents] assert query and documents and all(documents), \ f"query or documents is empty! query={query}, documents={len(documents)}" - + assert len(documents) < 500, \ + f"The input documents of Dashscope rerank model should not larger than 500!" # Using -1.0 as dummy scores nodes = [NodeWithScore(node=Node(text=doc), score=-1.0) for doc in documents] model_response.meta_data.update({ - "data": {"nodes": nodes, "query_str": query}, + "data": {"nodes": nodes, "query_str": query, "top_n": len(documents)}, "documents_map": {doc: idx for idx, doc in enumerate(documents)}, }) @@ -76,6 +77,8 @@ class LlamaIndexRankModel(BaseModel): Returns: ModelResponse: A response object encapsulating the ranked nodes. """ + self.model.top_n = model_response.meta_data["data"]["top_n"] + model_response.meta_data["data"].pop("top_n") model_response.raw = self.model.postprocess_nodes(**model_response.meta_data["data"]) async def _async_call(self, **kwargs) -> ModelResponse: diff --git a/memoryscope/storage/llama_index_es_memory_store.py b/memoryscope/storage/llama_index_es_memory_store.py index b7ce3b21..4a9d7f4f 100644 --- a/memoryscope/storage/llama_index_es_memory_store.py +++ b/memoryscope/storage/llama_index_es_memory_store.py @@ -26,7 +26,10 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): self.embedding_model: BaseModel = embedding_model self.es_store = SyncElasticsearchStore(index_name=index_name, es_url=es_url, - retrieval_strategy=_AsyncDenseVectorStrategy(hybrid=use_hybrid), + retrieval_strategy=_AsyncDenseVectorStrategy(hybrid=use_hybrid, + alpha=0.5), # weights of vector similarity, + # while the weights of BM25 is 1-alpha. + # when alpha=None, then rrf fusion is uesd. **kwargs) # TODO The llamaIndex utilizes some deprecated functions, hence langchain logs warning messages. By diff --git a/memoryscope/storage/llama_index_sync_elasticsearch.py b/memoryscope/storage/llama_index_sync_elasticsearch.py index 2c525920..507044a3 100644 --- a/memoryscope/storage/llama_index_sync_elasticsearch.py +++ b/memoryscope/storage/llama_index_sync_elasticsearch.py @@ -132,6 +132,21 @@ def _mode_must_match_retrieval_strategy( class _AsyncDenseVectorStrategy(AsyncDenseVectorStrategy): + + + def __init__( + self, + *, + distance: DistanceMetric = DistanceMetric.COSINE, + model_id: Optional[str] = None, + hybrid: bool = False, + rrf: Union[bool, Dict[str, Any]] = True, + text_field: Optional[str] = "text_field", + alpha: Optional[float] = None, + ): + super().__init__(distance=distance, model_id=model_id, hybrid=hybrid, rrf=rrf, text_field=text_field) + self.alpha = alpha + def _hybrid(self, query: str, knn: Dict[str, Any], filter: List[Dict[str, Any]], top_k: int) -> Dict[str, Any]: # Add a query to the knn query. # RRF is used to even the score from the knn query and text query @@ -155,18 +170,19 @@ class _AsyncDenseVectorStrategy(AsyncDenseVectorStrategy): "match": { self.text_field: { "query": query, + "boost": (1 - self.alpha) if self.alpha is not None else 1.0, } - } + }, } ], "filter": filter, - } + }, }, } - if isinstance(self.rrf, Dict): + if self.alpha is None and isinstance(self.rrf, Dict): query_body["rank"] = {"rrf": self.rrf} - elif isinstance(self.rrf, bool) and self.rrf is True: + elif self.alpha is None and isinstance(self.rrf, bool) and self.rrf is True: query_body["rank"] = {"rrf": {"window_size": top_k}} return query_body @@ -190,6 +206,7 @@ class _AsyncDenseVectorStrategy(AsyncDenseVectorStrategy): "field": vector_field, "k": k, "num_candidates": num_candidates, + "boost": self.alpha if self.alpha is not None else 1.0, } if query_vector is not None: @@ -676,6 +693,7 @@ class SyncElasticsearchStore(BasePydanticVectorStore): ): total_rank = sum(top_k_scores) top_k_scores = [rank for rank in top_k_scores] + print("top_k_scores:", top_k_scores) # top_k_scores = [(total_rank - rank) / total_rank for rank in top_k_scores] # top_k_scores = [total_rank - rank / total_rank for rank in top_k_scores]