From 064e8dda7022ec6c70d2a4a8f3f834e6cc6f69b3 Mon Sep 17 00:00:00 2001 From: ouguoyu Date: Wed, 20 Nov 2024 16:37:01 +0800 Subject: [PATCH] empty dir and file fixes --- examples/rag_milvus_store.py | 3 --- lazyllm/tools/rag/doc_impl.py | 2 +- lazyllm/tools/rag/rerank.py | 20 +++++++++++--------- 3 files changed, 12 insertions(+), 13 deletions(-) diff --git a/examples/rag_milvus_store.py b/examples/rag_milvus_store.py index 30b280bf..c2a23c9b 100644 --- a/examples/rag_milvus_store.py +++ b/examples/rag_milvus_store.py @@ -11,9 +11,6 @@ def __init__(self): self.root_dir = os.path.expanduser(os.path.join(config['home'], 'rag_for_ut')) self.rag_dir = os.path.join(self.root_dir, 'rag_master') os.makedirs(self.rag_dir, exist_ok=True) - # creates a dummy file for rag - with open(os.path.join(self.rag_dir, '_dummy.txt'), "wb") as fd: - fd.write(b'dsfjfasfkjdsfewifjewofjefiejw') self.store_file = os.path.join(self.root_dir, "milvus.db") def __del__(self): diff --git a/lazyllm/tools/rag/doc_impl.py b/lazyllm/tools/rag/doc_impl.py index 496d0680..5065fe36 100644 --- a/lazyllm/tools/rag/doc_impl.py +++ b/lazyllm/tools/rag/doc_impl.py @@ -255,7 +255,7 @@ def _add_files(self, input_files: List[str], ids: Optional[List[str]] = None, all_groups = self.store.all_groups() LOG.info(f"add_files: Trying to merge store with {all_groups}") for group in all_groups: - if not self.store.is_group_active(group): + if group != LAZY_ROOT_NAME and not self.store.is_group_active(group): continue # Duplicate group will be discarded automatically nodes = self._get_nodes(group, temp_store) diff --git a/lazyllm/tools/rag/rerank.py b/lazyllm/tools/rag/rerank.py index 43b97623..0e71aefb 100644 --- a/lazyllm/tools/rag/rerank.py +++ b/lazyllm/tools/rag/rerank.py @@ -95,15 +95,17 @@ def __init__(self, name: str = "ModuleReranker", target: Optional[str] = None, self._reranker = lazyllm.TrainableModule(self._kwargs['model']) def forward(self, nodes: List[DocNode], query: str = "") -> List[DocNode]: - docs = [node.get_text(metadata_mode=MetadataMode.EMBED) for node in nodes] - top_n = self._kwargs['topk'] if 'topk' in self._kwargs else len(docs) - if self._reranker._deploy_type == lazyllm.deploy.Infinity: - sorted_indices = self._reranker(query, documents=docs, top_n=top_n) - else: - inps = {'query': query, 'documents': docs, 'top_n': top_n} - sorted_indices = self._reranker(inps) - results = [nodes[i] for i in sorted_indices] - LOG.debug(f"Rerank use `{self._name}` and get nodes: {results}") + results = [] + if nodes: + docs = [node.get_text(metadata_mode=MetadataMode.EMBED) for node in nodes] + top_n = self._kwargs['topk'] if 'topk' in self._kwargs else len(docs) + if self._reranker._deploy_type == lazyllm.deploy.Infinity: + sorted_indices = self._reranker(query, documents=docs, top_n=top_n) + else: + inps = {'query': query, 'documents': docs, 'top_n': top_n} + sorted_indices = self._reranker(inps) + results = [nodes[i] for i in sorted_indices] + LOG.debug(f"Rerank use `{self._name}` and get nodes: {results}") return self._post_process(results) # User-defined similarity decorator