bug fix for EC-RAG

Signed-off-by: Zhu, Yongbo <[email protected]>
opea-project · Jan 16, 2025 · 37da801 · 37da801
1 parent b4269d6
commit 37da801
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/EdgeCraftRAG/edgecraftrag/components/retriever.py b/EdgeCraftRAG/edgecraftrag/components/retriever.py
@@ -79,6 +79,14 @@ def run(self, **kwargs) -> Any:
 
         return None
 
+    @model_serializer
+    def ser_model(self):
+        set = {
+            "idx": self.idx,
+            "retriever_type": self.comp_subtype,
+            "retrieve_topk": self.topk,
+        }
+        return set
 
 class SimpleBM25Retriever(BaseComponent):
     # The nodes parameter in BM25Retriever is not from index,
@@ -103,3 +111,12 @@ def run(self, **kwargs) -> Any:
                 return bm25_retr.retrieve(v)
 
         return None
+
+    @model_serializer
+    def ser_model(self):
+        set = {
+            "idx": self.idx,
+            "retriever_type": self.comp_subtype,
+            "retrieve_topk": self.topk,
+        }
+        return set