flashrag-dev 0.1.4.dev20250410__tar.gz → 0.1.4.dev20250424__tar.gz

{flashrag_dev-0.1.4.dev20250410 → flashrag_dev-0.1.4.dev20250424}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag_dev
-Version: 0.1.4.dev20250410
+Version: 0.1.4.dev20250424
 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.4.dev20250410 → flashrag_dev-0.1.4.dev20250424}/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.dev20250410 → flashrag_dev-0.1.4.dev20250424}/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.dev20250410 → flashrag_dev-0.1.4.dev20250424}/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.dev20250410 → flashrag_dev-0.1.4.dev20250424}/flashrag/retriever/encoder.py

@@ -24,7 +24,7 @@ class Encoder:
             Encodes a list of queries into embeddings.
     """
 
-    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16, instruction, silent=False):
+    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16=True, instruction="", silent=False):
         self.model_name = model_name
         self.model_path = model_path
         self.pooling_method = pooling_method
@@ -44,7 +44,9 @@ class Encoder:
         )
         inputs = {k: v.cuda() for k, v in inputs.items()}
 
-        if "T5" in type(self.model).__name__ or (isinstance(self.model, torch.nn.DataParallel) and "T5" in type(self.model.module).__name__):
+        if "T5" in type(self.model).__name__ or (
+            isinstance(self.model, torch.nn.DataParallel) and "T5" in type(self.model.module).__name__
+        ):
             # T5-based retrieval model
             decoder_input_ids = torch.zeros((inputs["input_ids"].shape[0], 1), dtype=torch.long).to(
                 inputs["input_ids"].device
@@ -54,11 +56,9 @@ class Encoder:
 
         else:
             output = self.model(**inputs, return_dict=True)
-            pooler_output = output.get('pooler_output', None)
-            last_hidden_state = output.get('last_hidden_state', None)
-            query_emb = pooling(
-                pooler_output, last_hidden_state, inputs["attention_mask"], self.pooling_method
-            )
+            pooler_output = output.get("pooler_output", None)
+            last_hidden_state = output.get("last_hidden_state", None)
+            query_emb = pooling(pooler_output, last_hidden_state, inputs["attention_mask"], self.pooling_method)
         if "dpr" not in self.model_name:
             query_emb = torch.nn.functional.normalize(query_emb, dim=-1)
         query_emb = query_emb.detach().cpu().numpy()
@@ -117,7 +117,11 @@ class STEncoder:
     def encode(self, query_list: Union[List[str], str], batch_size=64, is_query=True) -> np.ndarray:
         query_list = parse_query(self.model_name, query_list, self.instruction, is_query)
         query_emb = self.model.encode(
-            query_list, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=not self.silent
+            query_list,
+            batch_size=batch_size,
+            convert_to_numpy=True,
+            normalize_embeddings=True,
+            show_progress_bar=not self.silent,
         )
         query_emb = query_emb.astype(np.float32, order="C")
 
@@ -173,11 +177,10 @@ class ClipEncoder:
         # set model max length for model that not specified in config.json
         if self.processor is not None and self.processor.tokenizer.model_max_length > 100000:
             try:
-                model_max_length = config['text_config']['max_position_embeddings']
+                model_max_length = config["text_config"]["max_position_embeddings"]
             except:
                 model_max_length = 512
-            self.processor.tokenizer.model_max_length = model_max_length    
-
+            self.processor.tokenizer.model_max_length = model_max_length
 
     @torch.inference_mode()
     def single_batch_encode(self, query_list: Union[List[str], str], modal="image") -> np.ndarray:
@@ -207,7 +210,7 @@ class ClipEncoder:
     @torch.inference_mode()
     def encode_image(self, image_list: List) -> np.ndarray:
         # Each item in image_list: PIL Image, local path, or URL
-        if self.model_type == "CLIPModel" or self.model_type == 'ChineseCLIPModel':
+        if self.model_type == "CLIPModel" or self.model_type == "ChineseCLIPModel":
             # need handle image
             image_list = [parse_image(image) for image in image_list]
             inputs = self.processor(images=image_list, return_tensors="pt")
@@ -224,7 +227,7 @@ class ClipEncoder:
     @torch.inference_mode()
     def encode_text(self, text_list: List[str]) -> np.ndarray:
         # Each item in image_list: PIL Image, local path, or URL
-        if self.model_type == "CLIPModel" or self.model_type == 'ChineseCLIPModel':
+        if self.model_type == "CLIPModel" or self.model_type == "ChineseCLIPModel":
             inputs = self.processor(
                 text=text_list,
                 padding=True,

{flashrag_dev-0.1.4.dev20250410 → flashrag_dev-0.1.4.dev20250424}/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.dev20250424/flashrag/version.py

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

{flashrag_dev-0.1.4.dev20250410 → flashrag_dev-0.1.4.dev20250424}/flashrag_dev.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.4.dev20250410
+Version: 0.1.4.dev20250424
 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.4.dev20250410/flashrag/version.py

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