flashrag-dev 0.1.3.dev20241120__tar.gz → 0.1.3.dev20241122__tar.gz

{flashrag-dev-0.1.3.dev20241120/flashrag_dev.egg-info → flashrag-dev-0.1.3.dev20241122}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.3.dev20241120
+Version: 0.1.3.dev20241122
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou

{flashrag-dev-0.1.3.dev20241120 → flashrag-dev-0.1.3.dev20241122}/flashrag/dataset/dataset.py

@@ -176,7 +176,9 @@ class Dataset:
         save_data = [item.to_dict() for item in self.data]
         def custom_serializer(obj):
             if isinstance(obj, np.float32):  
-                return float(obj)           
+                return float(obj)      
+            if isinstance(obj, np.bool_):
+                return str(obj)     
             raise TypeError(f"Type {type(obj)} not serializable")
         with open(save_path, "w", encoding="utf-8") as f:
             json.dump(save_data, f, indent=4, default=custom_serializer)

{flashrag-dev-0.1.3.dev20241120 → flashrag-dev-0.1.3.dev20241122}/flashrag/generator/generator.py

@@ -176,14 +176,14 @@ class VLLMGenerator(BaseGenerator):
                 gpu_memory_utilization=gpu_memory_utilization,
                 enable_lora=True,
                 max_lora_rank=64,
-                max_logprobs=32016,
+                max_logprobs=130000,
             )
         else:
             self.model = LLM(
                 self.model_path,
                 tensor_parallel_size=tensor_parallel_size,
                 gpu_memory_utilization=gpu_memory_utilization,
-                max_logprobs=32016,
+                max_logprobs=130000,
             )
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
 

{flashrag-dev-0.1.3.dev20241120 → flashrag-dev-0.1.3.dev20241122}/flashrag/pipeline/active_pipeline.py

@@ -137,6 +137,7 @@ class SelfRAGPipeline(BasicPipeline):
         self.ret_tokens, self.rel_tokens, self.grd_tokens, self.ut_tokens = self.load_special_tokens(
             tokenizer, use_grounding=use_grounding, use_utility=use_utility
         )
+        self.vocab_size = tokenizer.vocab_size + len(tokenizer.added_tokens_decoder)
 
     def load_special_tokens(self, tokenizer, use_grounding, use_utility):
         ret_tokens = {token: tokenizer.convert_tokens_to_ids(token) for token in self.retrieval_tokens_names}
@@ -162,27 +163,25 @@ class SelfRAGPipeline(BasicPipeline):
         """Calculate whether a retrieve is required based on the output probability of
         the special token in the model"""
 
-        if self.mode != "always_retrieve":
+        if self.mode == "always_retrieve":
+            retrieval_flags = [True] * len(input_prompts)
+
+        elif self.mode == "no_retrieval":
+            retrieval_flags = [False] * len(input_prompts)
+
+        else:
             # result for total batch
-            all_pred_token_ids = []
             all_pred_text = []
             all_pred_log_probs = []
-            preds = self.generator.generate(input_prompts, return_raw_output=True, logprobs=32000)
+            # For vllm, requesting too many logprobes can seriously affect speed
+            # 20 probs is enough for calculate
+            preds = self.generator.generate(input_prompts, return_raw_output=True, logprobs=20, max_tokens=1, skip_special_tokens=False)
             for single_pred in preds:
-                pred_token_ids = single_pred.outputs[0].token_ids
                 pred_text = single_pred.outputs[0].text
                 pred_log_probs = single_pred.outputs[0].logprobs
-                all_pred_token_ids.append(pred_token_ids)
                 all_pred_text.append(pred_text)
                 all_pred_log_probs.append(pred_log_probs)
 
-        if self.mode == "always_retrieve":
-            retrieval_flags = [True] * len(input_prompts)
-
-        elif self.mode == "no_retrieval":
-            retrieval_flags = [False] * len(input_prompts)
-
-        else:
             retrieval_flags = []
             for idx, single_pred in enumerate(preds):
                 if self.threshold is not None:
@@ -190,8 +189,9 @@ class SelfRAGPipeline(BasicPipeline):
                     for tok, tok_id in self.ret_tokens.items():
                         if tok_id not in all_pred_log_probs[idx][0]:
                             score_dict[tok] = -100
-                        prob = all_pred_log_probs[idx][0][tok_id].logprob
-                        score_dict[tok] = float(prob)
+                        else:
+                            prob = all_pred_log_probs[idx][0][tok_id].logprob
+                            score_dict[tok] = np.exp(prob)
                     do_retrieve = (
                         score_dict["[Retrieval]"] / (score_dict["[Retrieval]"] + score_dict["[No Retrieval]"])
                         > self.threshold
@@ -417,7 +417,7 @@ class SelfRAGPipeline(BasicPipeline):
                         else:
                             aug_prompts = [prompt + prev_generation]
 
-                        item_pred = self.generator.generate(aug_prompts, return_raw_output=True)
+                        item_pred = self.generator.generate(aug_prompts, return_raw_output=True, logprobs=5)
                         _, preds, scores, overall_score_dict = self.critic_preds(item_pred)
 
                         for i, (pred, p_score) in enumerate(zip(preds, scores)):
@@ -602,15 +602,17 @@ class SelfRAGPipeline(BasicPipeline):
 
         return dataset
 
-    def run(self, dataset, do_eval=True, pred_process_fun=None, batch_size=256, long_form=False):
-        all_dataset_list = []
+    def run(self, dataset, do_eval=True, pred_process_fun=None, long_form=False):
         run_func = self.run_batch_pred_long_form if long_form else self.run_batch_pred
-        # to avoid oom
-        for batch_dataset in tqdm(get_batch_dataset(dataset, batch_size=batch_size), desc="Batch dataset: "):
-            batch_dataset = run_func(batch_dataset)
-            all_dataset_list.append(batch_dataset)
-        dataset = merge_batch_dataset(all_dataset_list)
-
+        
+        # # to avoid oom, split the total dataset into small batches
+        # all_dataset_list = []
+        # for batch_dataset in tqdm(get_batch_dataset(dataset, batch_size=batch_size), desc="Batch dataset: "):
+        #     batch_dataset = run_func(batch_dataset)
+        #     all_dataset_list.append(batch_dataset)
+        # dataset = merge_batch_dataset(all_dataset_list)
+
+        dataset = run_func(dataset)
         dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
         return dataset
 
@@ -645,7 +647,7 @@ class SelfRAGPipeline(BasicPipeline):
             item.update_output("prompt", prompt_list)
             all_input_list += prompt_list
 
-        batch_pred = self.generator.generate(all_input_list, return_raw_output=True, logprobs=32016)
+        batch_pred = self.generator.generate(all_input_list, return_raw_output=True, logprobs=5)
 
         # parse output based on retrieval flag
         pred_idx = 0

{flashrag-dev-0.1.3.dev20241120 → flashrag-dev-0.1.3.dev20241122}/flashrag/retriever/index_builder.py

@@ -165,14 +165,19 @@ class Index_Builder:
         """Building BM25 index based on bm25s library."""
 
         import bm25s
+        import Stemmer
         
         self.save_dir = os.path.join(self.save_dir, 'bm25')
         os.makedirs(self.save_dir, exist_ok=True)
 
         corpus = datasets.load_dataset("json", data_files=self.corpus_path, split="train")
         corpus_text = corpus['contents']
+        stemmer = Stemmer.Stemmer('english')
+        tokenizer = bm25s.tokenization.Tokenizer(stemmer=stemmer)
+        corpus_tokens = tokenizer.tokenize(corpus_text, return_as='tuple')
+
         retriever = bm25s.BM25(corpus=corpus, backend='numba')
-        retriever.index(corpus_text)
+        retriever.index(corpus_tokens)
         retriever.save(self.save_dir,corpus=corpus)
 
         print("Finish!")

flashrag-dev-0.1.3.dev20241122/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.3dev20241122"

{flashrag-dev-0.1.3.dev20241120 → flashrag-dev-0.1.3.dev20241122/flashrag_dev.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.3.dev20241120
+Version: 0.1.3.dev20241122
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou

{flashrag-dev-0.1.3.dev20241120 → flashrag-dev-0.1.3.dev20241122}/flashrag_dev.egg-info/requires.txt

@@ -7,12 +7,12 @@ peft
 PyYAML
 rank_bm25
 rouge
-spacy>=3.6.1
+spacy==3.6.1
 tiktoken
 torch
 tqdm
 transformers>=4.40.0
-bm25s[full]
+bm25s[core]==0.2.0
 fschat
 streamlit
 
@@ -26,12 +26,12 @@ peft
 PyYAML
 rank_bm25
 rouge
-spacy>=3.6.1
+spacy==3.6.1
 tiktoken
 torch
 tqdm
 transformers>=4.40.0
-bm25s[full]
+bm25s[core]==0.2.0
 fschat
 streamlit
 
@@ -45,12 +45,12 @@ peft
 PyYAML
 rank_bm25
 rouge
-spacy>=3.6.1
+spacy==3.6.1
 tiktoken
 torch
 tqdm
 transformers>=4.40.0
-bm25s[full]
+bm25s[core]==0.2.0
 fschat
 streamlit
 pyserini

flashrag-dev-0.1.3.dev20241120/flashrag/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.3dev20241120"