flashrag-dev 0.1.4.dev20250403__tar.gz → 0.1.4.dev20250423__tar.gz

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag_dev
-Version: 0.1.4.dev20250403
+Version: 0.1.4.dev20250423
 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
@@ -94,6 +94,7 @@ FlashRAG is still under development and there are many issues and room for impro
 - [ ] Inlcude more RAG approaches
 - [ ] Add more evaluation metrics (e.g., Unieval, name-entity F1) and benchmarks (e.g., RGB benchmark)
 - [ ] Enhance code adaptability and readability
+- [ ] Add support for api-based retriever (vllm server)
 
 ## :page_with_curl: Changelog
 [25/03/21] We have added support for **Reasoning Pipeline**, which is a new paradigm that combines reasoning ability and retrieval, representing work that includes [Search-o1](https://github.com/sunnynexus/Search-o1), [R1-Searcher](https://github.com/SsmallSong/R1-Searcher), [ReSearch](https://github.com/Agent-RL/ReSearch). We evaluate the performance of the pipeline on various RAG benchmarks, it can achieve F1 scores close to 60 on multi hop inference datasets such as Hotpotqe. See it in [**result table**](#robot-supporting-methods).
@@ -108,6 +109,8 @@ FlashRAG is still under development and there are many issues and room for impro
 
 [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/original_docs/multi_retriever_usage.md).
 
+[25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/original_docs/chunk-doc-corpus.md).
+
 [24/10/21] We have released a version based on the Paddle framework that supports Chinese hardware platforms. Please refer to [FlashRAG Paddle](https://github.com/RUC-NLPIR/FlashRAG-Paddle) for details.
 
 [24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/README.md

@@ -77,6 +77,7 @@ FlashRAG is still under development and there are many issues and room for impro
 - [ ] Inlcude more RAG approaches
 - [ ] Add more evaluation metrics (e.g., Unieval, name-entity F1) and benchmarks (e.g., RGB benchmark)
 - [ ] Enhance code adaptability and readability
+- [ ] Add support for api-based retriever (vllm server)
 
 ## :page_with_curl: Changelog
 [25/03/21] We have added support for **Reasoning Pipeline**, which is a new paradigm that combines reasoning ability and retrieval, representing work that includes [Search-o1](https://github.com/sunnynexus/Search-o1), [R1-Searcher](https://github.com/SsmallSong/R1-Searcher), [ReSearch](https://github.com/Agent-RL/ReSearch). We evaluate the performance of the pipeline on various RAG benchmarks, it can achieve F1 scores close to 60 on multi hop inference datasets such as Hotpotqe. See it in [**result table**](#robot-supporting-methods).
@@ -91,6 +92,8 @@ FlashRAG is still under development and there are many issues and room for impro
 
 [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/original_docs/multi_retriever_usage.md).
 
+[25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/original_docs/chunk-doc-corpus.md).
+
 [24/10/21] We have released a version based on the Paddle framework that supports Chinese hardware platforms. Please refer to [FlashRAG Paddle](https://github.com/RUC-NLPIR/FlashRAG-Paddle) for details.
 
 [24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/flashrag/generator/openai_generator.py

@@ -1,5 +1,5 @@
 import os
-from typing import List
+from typing import List, Union
 from copy import deepcopy
 import warnings
 from tqdm import tqdm
@@ -63,16 +63,26 @@ class OpenaiGenerator:
     def update_additional_setting(self):
         pass
     
-    async def _get_response(self, messages: List, **params):
-        response = await self.client.chat.completions.create(
-            model=self.model_name, messages=messages, **params
-        )
-        if not response.choices:
-            raise ValueError("No choices returned from API.")
-        return response.choices[0]
+    async def _get_response(self, messages: Union[list, str], mode: str = 'chat', **params):
+        print(messages)
+        print("----asd-asd-a--------")
+        if mode == 'chat':
+            response = await self.client.chat.completions.create(
+                model=self.model_name, messages=messages, **params
+            )
+            if not response.choices:
+                raise ValueError("No choices returned from API.")
+            return response.choices[0]
+        else:
+            response = await self.client.completions.create(
+                model=self.model_name, prompt=messages, **params
+            )
+            if not response.choices:
+                raise ValueError("No choices returned from API.")
+            return response.choices[0]
 
-    async def _get_batch_response(self, input_list: List[List], batch_size, **params):
-        tasks = [self._get_response(messages, **params) for messages in input_list]
+    async def _get_batch_response(self, input_list: List[List], batch_size, mode, **params):
+        tasks = [self._get_response(messages, mode, **params) for messages in input_list]
         all_results = []
         for idx in tqdm(range(0, len(tasks), batch_size), desc="Generation process: "):
             batch_tasks = tasks[idx: idx + batch_size]
@@ -85,6 +95,10 @@ class OpenaiGenerator:
             input_list = [[input_list]]
         elif isinstance(input_list[0], dict):
             input_list = [input_list]
+        if isinstance(input_list[0], list):
+            mode = 'chat'
+        else:
+            mode = 'completion'
 
         if batch_size is None:
             batch_size = self.batch_size
@@ -107,15 +121,26 @@ class OpenaiGenerator:
             warnings.warn("Set logprobs to True to get generation scores.")
 
 
-        results = await self._get_batch_response(input_list, batch_size, **generation_params)
+        results = await self._get_batch_response(input_list, batch_size, mode, **generation_params)
 
         response_texts = []
         scores = []
         for res in results:
-            response_texts.append(res.message.content)
+            if mode == 'chat':
+                text = res.message.content
+            else:
+                text = res.text
+            if 'stop' in generation_params and not any([text.endswith(stop_str) for stop_str in generation_params['stop']]):
+                # current openai api not support including stop, so need to add the stop token
+                text += generation_params['stop'][0]
+            response_texts.append(text)
             if return_scores:
-                score = np.exp([item.logprob for item in res.logprobs.content])
-                scores.append(score)
+                try:
+                    score = np.exp([item.logprob for item in res.logprobs.content])
+                    scores.append(score)
+                except:
+                    warnings.warn('Fail to get logprobs in openai generation!')
+                    scores.append(None)
         return (response_texts, scores) if return_scores else response_texts
 
     # ----------------- 同步包装接口 -----------------

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/flashrag/pipeline/active_pipeline.py

@@ -931,12 +931,17 @@ class IRCOTPipeline(BasicPipeline):
     ):
         # if not provide prompt template, use default template provided by IRCOT
         if prompt_template is None:
+            if config['framework'] == 'openai':
+                enable_chat = True
+            else:
+                enable_chat = False
+            
             prompt_template = PromptTemplate(
                 config=config,
                 system_prompt=f"{self.IRCOT_INSTRUCTION}\n\n{self.IRCOT_EXAMPLE}",
                 user_prompt="{reference}Question: {question}\nThought:",
                 reference_template="Wikipedia Title: {title}\n{text}\n\n",
-                enable_chat=False,
+                enable_chat=enable_chat,
             )
 
         super().__init__(config, prompt_template)
@@ -957,7 +962,6 @@ class IRCOTPipeline(BasicPipeline):
         questions = [item.question for item in items]
         retrieval_results, scoress = self.retriever.batch_search(questions, return_score=True)
         for retrieval_result, scores in zip(retrieval_results,scoress):   
-            
             doc2score = {doc_item['id']: score for doc_item, score in zip(retrieval_result, scores)}
             id2doc = {doc_item['id']: doc_item for doc_item in retrieval_result}
             batch_retrieval_results.append(retrieval_result)
@@ -976,6 +980,7 @@ class IRCOTPipeline(BasicPipeline):
                 )
                 for item_id in active_item_ids
             ]
+            
 
             # Batch generation for active items
             new_thoughts_batch = self.generator.generate(input_prompts, stop=['.', '\n'])
@@ -988,15 +993,14 @@ class IRCOTPipeline(BasicPipeline):
                 
                 # Check for termination condition
                 # Store intermediate outputs
-                if "So the answer is:" in new_thought:
-                    items[item_id].update_output(
+                items[item_id].update_output(
                         f'intermediate_output_iter{iter_num}', 
                         {
                             'input_prompt': input_prompts[idx],
                             'new_thought': new_thought,
                         },
                     )
-                else:
+                if "So the answer is:" not in new_thought:
                     new_active_item_ids.append(item_id)
 
             # Update active item IDs for the next iteration

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/flashrag/prompt/base_prompt.py

@@ -58,23 +58,29 @@ class PromptTemplate:
 
     def truncate_prompt(self, prompt):
         if self.is_openai:
-            truncated_messages = []
-            total_tokens = 0
-            assert isinstance(prompt, list)
-            for message in prompt:
-                role_content = message['content']
-                encoded_message = self.tokenizer.encode(role_content)
-
-                if total_tokens + len(encoded_message) <= self.max_input_len:
-                    truncated_messages.append(message)
-                    total_tokens += len(encoded_message)
-                else:
-                    print(f"The input text length is greater than the maximum length ({total_tokens + len(encoded_message)} > {self.max_input_len}) and has been truncated!")
-                    remaining_tokens = self.max_input_len - total_tokens
-                    truncated_message = self.encoding.decode(encoded_message[:remaining_tokens])
-                    message['content'] = truncated_message
-                    truncated_messages.append(message)
-                    break
+            if self.enable_chat:
+                truncated_messages = []
+                total_tokens = 0
+                assert isinstance(prompt, list)
+                for message in prompt:
+                    role_content = message['content']
+                    encoded_message = self.tokenizer.encode(role_content)
+
+                    if total_tokens + len(encoded_message) <= self.max_input_len:
+                        truncated_messages.append(message)
+                        total_tokens += len(encoded_message)
+                    else:
+                        print(f"The input text length is greater than the maximum length ({total_tokens + len(encoded_message)} > {self.max_input_len}) and has been truncated!")
+                        remaining_tokens = self.max_input_len - total_tokens
+                        truncated_message = self.encoding.decode(encoded_message[:remaining_tokens])
+                        message['content'] = truncated_message
+                        truncated_messages.append(message)
+                        break
+            else:
+                assert isinstance(prompt, str)
+                tokenized_prompt = self.tokenizer.encode(prompt,allowed_special={'<|endoftext|>'})
+                half = int(self.max_input_len / 2)
+                truncated_messages = self.tokenizer.decode(tokenized_prompt[:half]) + self.tokenizer.decode(tokenized_prompt[-half:])
 
             return truncated_messages
 
@@ -134,8 +140,15 @@ class PromptTemplate:
         else:
             input = "\n\n".join([prompt for prompt in [system_prompt, user_prompt] if prompt != ""])
 
-        if previous_gen is not None and previous_gen not in ["", " "] and self.is_openai is False:
-            input += previous_gen
+        if previous_gen is not None and previous_gen not in ["", " "]:
+            if self.is_openai:
+                if self.enable_chat:
+                    input.append({"role": 'assistant', 'content': previous_gen})
+                else:    
+                    input += f'<|endoftext|>{previous_gen}'
+                
+            else:
+                input += previous_gen
 
         return self.truncate_prompt(input)
 

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/flashrag/retriever/reranker.py

@@ -44,7 +44,7 @@ class BaseReranker:
             doc_list = [doc_list]
 
         assert len(query_list) == len(doc_list)
-        if topk < min([len(docs) for docs in doc_list]):
+        if topk > min([len(docs) for docs in doc_list]):
             warnings.warn("The number of doc returned by the retriever is less than the topk.")
 
         # get doc contents

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/flashrag/retriever/retriever.py

@@ -242,10 +242,7 @@ class BM25Retriever(BaseTextRetriever):
                 else:
                     self.corpus = corpus
             self.max_process_num = 8
-
-            is_zh = judge_zh(self.corpus[0]["contents"])
-            if is_zh:
-                self.searcher.set_language("zh")
+               
         elif self.backend == "bm25s":
             import Stemmer
             import bm25s
@@ -278,6 +275,9 @@ class BM25Retriever(BaseTextRetriever):
         if num is None:
             num = self.topk
         if self.backend == "pyserini":
+            is_zh = judge_zh(query)
+            if is_zh:
+                self.searcher.set_language("zh")
             hits = self.searcher.search(query, num)
             if len(hits) < 1:
                 if return_score:
@@ -295,11 +295,12 @@ class BM25Retriever(BaseTextRetriever):
                 all_contents = [json.loads(self.searcher.doc(hit.docid).raw())["contents"] for hit in hits]
                 results = [
                     {
+                        "id": hit.docid, 
                         "title": content.split("\n")[0].strip('"'),
                         "text": "\n".join(content.split("\n")[1:]),
                         "contents": content,
                     }
-                    for content in all_contents
+                    for content, hit in zip(all_contents, hits)
                 ]
             else:
                 results = load_docs(self.corpus, [hit.docid for hit in hits])

flashrag_dev-0.1.4.dev20250423/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.4dev20250423"

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/flashrag_dev.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.4.dev20250403
+Version: 0.1.4.dev20250423
 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
@@ -94,6 +94,7 @@ FlashRAG is still under development and there are many issues and room for impro
 - [ ] Inlcude more RAG approaches
 - [ ] Add more evaluation metrics (e.g., Unieval, name-entity F1) and benchmarks (e.g., RGB benchmark)
 - [ ] Enhance code adaptability and readability
+- [ ] Add support for api-based retriever (vllm server)
 
 ## :page_with_curl: Changelog
 [25/03/21] We have added support for **Reasoning Pipeline**, which is a new paradigm that combines reasoning ability and retrieval, representing work that includes [Search-o1](https://github.com/sunnynexus/Search-o1), [R1-Searcher](https://github.com/SsmallSong/R1-Searcher), [ReSearch](https://github.com/Agent-RL/ReSearch). We evaluate the performance of the pipeline on various RAG benchmarks, it can achieve F1 scores close to 60 on multi hop inference datasets such as Hotpotqe. See it in [**result table**](#robot-supporting-methods).
@@ -108,6 +109,8 @@ FlashRAG is still under development and there are many issues and room for impro
 
 [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/original_docs/multi_retriever_usage.md).
 
+[25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/original_docs/chunk-doc-corpus.md).
+
 [24/10/21] We have released a version based on the Paddle framework that supports Chinese hardware platforms. Please refer to [FlashRAG Paddle](https://github.com/RUC-NLPIR/FlashRAG-Paddle) for details.
 
 [24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.

{flashrag_dev-0.1.4.dev20250403 → flashrag_dev-0.1.4.dev20250423}/flashrag_dev.egg-info/requires.txt

@@ -16,7 +16,7 @@ transformers>=4.40.0
 bm25s[core]==0.2.0
 fschat
 streamlit
-chonkie>=0.4.0
+chonkie<1.1.0,>=1.0.2
 gradio>=5.0.0
 rouge-chinese
 jieba
@@ -40,7 +40,7 @@ transformers>=4.40.0
 bm25s[core]==0.2.0
 fschat
 streamlit
-chonkie>=0.4.0
+chonkie<1.1.0,>=1.0.2
 gradio>=5.0.0
 rouge-chinese
 jieba
@@ -64,7 +64,7 @@ transformers>=4.40.0
 bm25s[core]==0.2.0
 fschat
 streamlit
-chonkie>=0.4.0
+chonkie<1.1.0,>=1.0.2
 gradio>=5.0.0
 rouge-chinese
 jieba

flashrag_dev-0.1.4.dev20250403/flashrag/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.4dev20250403"