[WIP] New training code with improved rag and parallel tool call

examples/vllm_multiturn/config/tool_config/search_tool_config.yaml

@@ -2,9 +2,9 @@ tools:
   - class_name: verl.tools.janv2_tool.web_search_tool.WebSearchTool
     config:
       type: native
-      rag_server_url: "http://10.220.108.31:3030"
+      rag_server_url: "http://localhost:3030"
       num_results: 10
-      topk_retrieval: 30
+      topk_retrieval: 200
       num_workers: 64
       rate_limit: 100000
       timeout: 600
@@ -29,7 +29,7 @@ tools:
   - class_name: verl.tools.janv2_tool.scrape_tool.ScrapeTool
     config:
       type: native
-      rag_server_url: "http://10.220.108.31:3030"
+      rag_server_url: "http://localhost:3030"
       num_workers: 50
       rate_limit: 100000
       timeout: 600

examples/vllm_multiturn/rag_setup/bm25_index.py

@@ -1,7 +1,12 @@
 import bm25s
 import json, re
+import gzip
+import os
 import datasets
-import Stemmer 
+import Stemmer
+import logging
+
+logger = logging.getLogger(__name__)
 
 def load_corpus(corpus_path: str):
     """Load corpus using datasets library"""
@@ -20,35 +25,72 @@ def load_docs(corpus, doc_idxs):
     return results
 
 class BM25RetrieverLunce:
-    def __init__(self, corpus_path: str):
+    def __init__(self, corpus_path_or_corpus, is_corpus=False, cache_dir=None):
+        """
+        Args:
+            corpus_path_or_corpus: Either a path string or a pre-loaded corpus dataset
+            is_corpus: If True, corpus_path_or_corpus is a pre-loaded corpus
+            cache_dir: Directory to save/load BM25 index. If None, uses corpus_path + "_bm25_cache"
+        """
+        if is_corpus:
+            self.corpus = corpus_path_or_corpus
+            self.cache_dir = cache_dir
+        else:
+            logger.info("BM25: Loading corpus...")
+            self.corpus = load_corpus(corpus_path=corpus_path_or_corpus)
+            # Default cache dir based on corpus path
+            if cache_dir is None:
+                self.cache_dir = corpus_path_or_corpus + "_bm25_cache"
+            else:
+                self.cache_dir = cache_dir
 
-        self.retriever = self._build_index(corpus_path)
-        self.corpus = load_corpus(corpus_path=corpus_path)
-
-    def _build_index(self, corpus_path):
-        with open(corpus_path,"r") as file:
-            lines = file.readlines()
-        self.raw_data = []
-        for line in lines:
-            try:
-                data = json.loads(line)
-                self.raw_data.append(data)
-            except:
-                print(f"error when loading: {data}")
-        corpus = [re.sub(r'[^\w\s]', '', data["contents"]) for data  in self.raw_data] 
         self.stemmer = Stemmer.Stemmer("english")
-        retriever =  bm25s.BM25() #corpus=corpus
-        retriever.index(bm25s.tokenize(corpus, stopwords="en", stemmer=self.stemmer))
+        self.retriever = self._load_or_build_index()
+
+    def _load_or_build_index(self):
+        """Load index from cache if exists, otherwise build and save"""
+        if self.cache_dir and os.path.exists(self.cache_dir):
+            logger.info(f"BM25: Loading cached index from {self.cache_dir}...")
+            retriever = bm25s.BM25.load(self.cache_dir, load_corpus=False)
+            logger.info("BM25: Cached index loaded successfully!")
+            return retriever
+        else:
+            logger.info(f"BM25: No cached index found, building new index...")
+            retriever = self._build_index()
+
+            # Save to cache
+            if self.cache_dir:
+                logger.info(f"BM25: Saving index to {self.cache_dir}...")
+                os.makedirs(self.cache_dir, exist_ok=True)
+                retriever.save(self.cache_dir)
+                logger.info("BM25: Index saved to cache!")
+
+            return retriever
+
+    def _build_index(self):
+        logger.info(f"BM25: Building index for {len(self.corpus)} documents...")
+
+        # Extract texts directly from corpus (more efficient)
+        corpus_texts = [re.sub(r'[^\w\s]', '', doc["contents"]) for doc in self.corpus]
+
+        logger.info("BM25: Tokenizing corpus...")
+        tokens = bm25s.tokenize(corpus_texts, stopwords="en", stemmer=self.stemmer)
+
+        logger.info("BM25: Indexing tokens...")
+        retriever = bm25s.BM25()
+        retriever.index(tokens)
+
+        logger.info("BM25: Index built successfully!")
         return retriever
 
-    def _search(self,query: str, num: int):
+    def _search(self, query: str, num: int):
         results, scores = self.retriever.retrieve(bm25s.tokenize(query, stopwords="en", stemmer=self.stemmer), k=num)
         return results[0], scores[0]
 
 
 
 if __name__ == "__main__":
-    bm25_ = BM25Retriever("/mnt/nas/alex/deep-research/src/rag_setup/data/corpus/corpus.jsonl")
+    bm25_ = BM25RetrieverLunce("/mnt/nas/alex/deep-research/src/rag_setup/data/corpus/corpus.jsonl")
     print(bm25_._search("Mc Donald", 5))
     result = bm25_._search(" Donald", 5)
     print(load_docs(bm25_.corpus, result[0][0]))