flashrag-dev 0.1.4.dev20241203__tar.gz → 0.1.4.dev20241210__tar.gz

{flashrag-dev-0.1.4.dev20241203/flashrag_dev.egg-info → flashrag-dev-0.1.4.dev20241210}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.4.dev20241203
+Version: 0.1.4.dev20241210
 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.dev20241203 → flashrag-dev-0.1.4.dev20241210}/flashrag/config/config.py

@@ -1,5 +1,6 @@
 import re
 import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 import yaml
 import random
 import datetime
@@ -102,13 +103,19 @@ class Config:
         gpu_id = self.final_config["gpu_id"]
         if gpu_id is not None:
             os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+        try:
+            # import pynvml 
+            # pynvml.nvmlInit()
+            # gpu_num = pynvml.nvmlDeviceGetCount()
             import torch
-
-            self.final_config["device"] = torch.device("cuda")
+            gpu_num = torch.cuda.device_count()
+        except:
+            gpu_num = 0
+        self.final_config['gpu_num'] = gpu_num
+        if gpu_num > 0:
+            self.final_config["device"] = "cuda"
         else:
-            import torch
-
-            self.final_config["device"] = torch.device("cpu")
+            self.final_config['device'] = 'cpu'
 
     def _set_additional_key(self):
         # set dataset

{flashrag-dev-0.1.4.dev20241203 → flashrag-dev-0.1.4.dev20241210}/flashrag/generator/generator.py

@@ -4,7 +4,6 @@ import warnings
 from tqdm import tqdm
 from tqdm.auto import trange
 import numpy as np
-import torch
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
@@ -25,7 +24,7 @@ class BaseGenerator:
         self.max_input_len = config["generator_max_input_len"]
         self.batch_size = config["generator_batch_size"]
         self.device = config["device"]
-        self.gpu_num = torch.cuda.device_count()
+        self.gpu_num = config['gpu_num']
         self.config = config
         self.generation_params = config["generation_params"]
 
@@ -64,6 +63,7 @@ class EncoderDecoderGenerator(BaseGenerator):
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
 
     def encode_passages(self, batch_text_passages: List[List[str]]):
+        import torch
         # need size: [batch_size, passage_num, passage_len]
         passage_ids, passage_masks = [], []
         for text_passages in batch_text_passages:
@@ -81,7 +81,6 @@ class EncoderDecoderGenerator(BaseGenerator):
         passage_masks = torch.cat(passage_masks, dim=0)
         return passage_ids, passage_masks.bool()
 
-    @torch.inference_mode(mode=True)
     def generate(self, input_list: List, batch_size=None, **params):
         if isinstance(input_list, str):
             input_list = [input_list]
@@ -128,15 +127,17 @@ class EncoderDecoderGenerator(BaseGenerator):
                 ).to(self.device)
 
             # TODO: multi-gpu inference
-            if self.fid:
-                if 'max_new_tokens' in generation_params:
-                    max_new_tokens = generation_params.pop('max_new_tokens')
+            import torch
+            with torch.inference_mode():
+                if self.fid:
+                    if 'max_new_tokens' in generation_params:
+                        max_new_tokens = generation_params.pop('max_new_tokens')
+                    else:
+                        max_new_tokens = 32
+
+                    outputs = self.model.generate(**inputs, max_new_tokens=max_new_tokens, pad_token_id=self.tokenizer.pad_token_id, decoder_start_token_id=self.tokenizer.pad_token_id)
                 else:
-                    max_new_tokens = 32
-            
-                outputs = self.model.generate(**inputs, max_new_tokens=max_new_tokens, pad_token_id=self.tokenizer.pad_token_id, decoder_start_token_id=self.tokenizer.pad_token_id)
-            else:
-                outputs = self.model.generate(**inputs, **generation_params)
+                    outputs = self.model.generate(**inputs, **generation_params)
             outputs = self.tokenizer.batch_decode(
                 outputs,
                 skip_special_tokens=True,
@@ -153,9 +154,8 @@ class VLLMGenerator(BaseGenerator):
 
     def __init__(self, config):
         super().__init__(config)
-
+        
         from vllm import LLM
-
         if "gpu_memory_utilization" not in config:
             gpu_memory_utilization = 0.85
         else:
@@ -187,7 +187,6 @@ class VLLMGenerator(BaseGenerator):
             )
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
 
-    @torch.inference_mode(mode=True)
     def generate(
         self,
         input_list: List[str],
@@ -283,6 +282,7 @@ class HFCausalLMGenerator(BaseGenerator):
         return model, tokenizer
 
     def add_new_tokens(self, token_embedding_path, token_name_func=lambda idx: f"[ref{idx+1}]"):
+        import torch
         del self.model
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_path,
@@ -318,7 +318,6 @@ class HFCausalLMGenerator(BaseGenerator):
         self.model.eval()
         self.model.cuda()
 
-    @torch.inference_mode(mode=True)
     def generate(
         self,
         input_list: List[str],
@@ -370,28 +369,30 @@ class HFCausalLMGenerator(BaseGenerator):
         generated_token_ids = []
         generated_token_logits = []
 
+        import torch
         for idx in trange(0, len(input_list), batch_size, desc="Generation process: "):
-            torch.cuda.empty_cache()
-            batched_prompts = input_list[idx : idx + batch_size]
-            inputs = self.tokenizer(
-                batched_prompts,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=self.max_input_len,
-            ).to(self.model.device)
-            outputs = self.model.generate(
-                **inputs,
-                output_scores=True,
-                return_dict_in_generate=True,
-                **generation_params,
-            )
+            with torch.inference_mode():
+                torch.cuda.empty_cache()
+                batched_prompts = input_list[idx : idx + batch_size]
+                inputs = self.tokenizer(
+                    batched_prompts,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=self.max_input_len,
+                ).to(self.model.device)
+                outputs = self.model.generate(
+                    **inputs,
+                    output_scores=True,
+                    return_dict_in_generate=True,
+                    **generation_params,
+                )
 
-            generated_ids = outputs.sequences
-            logits = torch.stack(outputs.scores, dim=1).softmax(-1)
-            generated_ids = generated_ids[:, inputs["input_ids"].shape[-1] :]
-            gen_score = torch.gather(logits, 2, generated_ids[:, :, None]).squeeze(-1).cpu().tolist()
-            scores.extend(gen_score)
+                generated_ids = outputs.sequences
+                logits = torch.stack(outputs.scores, dim=1).softmax(-1)
+                generated_ids = generated_ids[:, inputs["input_ids"].shape[-1] :]
+                gen_score = torch.gather(logits, 2, generated_ids[:, :, None]).squeeze(-1).cpu().tolist()
+                scores.extend(gen_score)
 
             # get additinoal info
             if return_dict:
@@ -477,13 +478,14 @@ class HFCausalLMGenerator(BaseGenerator):
         else:
             return responses
 
-    @torch.inference_mode(mode=True)
+
     def cal_gen_probs(self, prev, next):
+        import torch
         input_ids = self.tokenizer.encode(prev, add_special_tokens=False)
         target_ids = self.tokenizer.encode(next, add_special_tokens=False)
         context_ids = input_ids + target_ids
         context_tensor = torch.tensor([context_ids]).to(self.device)
-        with torch.no_grad():
+        with torch.inference_mode():
             outputs = self.model(context_tensor)
             logits = outputs.logits
             logits = logits[0, len(input_ids) - 1 : len(context_ids) - 1, :]
@@ -505,6 +507,7 @@ class FastChatGenerator(HFCausalLMGenerator):
 
         def get_gpu_memory(max_gpus=None):
             """Get available memory for each GPU."""
+            import torch
             gpu_memory = []
             num_gpus = torch.cuda.device_count() if max_gpus is None else min(max_gpus, torch.cuda.device_count())
             for gpu_id in range(num_gpus):
@@ -525,6 +528,8 @@ class FastChatGenerator(HFCausalLMGenerator):
             else:
                 gpu_memory_utilization = self.config["gpu_memory_utilization"]
             max_gpu_memory = None
+            import torch
+            self.gpu_num = torch.cuda.device_count()
             if self.gpu_num != 1:
                 available_gpu_memory = get_gpu_memory(self.gpu_num)
                 max_gpu_memory = str(int(min(available_gpu_memory) * gpu_memory_utilization)) + "GiB"

{flashrag-dev-0.1.4.dev20241203 → flashrag-dev-0.1.4.dev20241210}/flashrag/pipeline/__init__.py

@@ -1,3 +1,3 @@
 from flashrag.pipeline.pipeline import *
 from flashrag.pipeline.branching_pipeline import REPLUGPipeline, SuRePipeline
-from flashrag.pipeline.active_pipeline import IterativePipeline, SelfRAGPipeline, FLAREPipeline, SelfAskPipeline, IRCOTPipeline
+from flashrag.pipeline.active_pipeline import IterativePipeline, SelfRAGPipeline, FLAREPipeline, SelfAskPipeline, IRCOTPipeline, RQRAGPipeline

{flashrag-dev-0.1.4.dev20241203 → flashrag-dev-0.1.4.dev20241210}/flashrag/pipeline/active_pipeline.py

@@ -1,5 +1,7 @@
 import re
 from tqdm import tqdm
+from typing import List, Tuple
+import math
 import numpy as np
 from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
 from flashrag.utils import get_retriever, get_generator, selfask_pred_parse, ircot_pred_parse
@@ -1014,3 +1016,220 @@ class IRCOTPipeline(BasicPipeline):
 
         dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
         return dataset
+
+class RQRAGPipeline(BasicPipeline):
+    expand_on_tokens = [
+        "[S_Rewritten_Query]",
+        "[S_Decomposed_Query]",
+        "[S_Disambiguated_Query]",
+        "[A_Response]"
+    ]
+    
+    system_prompt = {
+        "qa": "Given a question that requires multi-hop reasoning, you need to decompose the question and answer based on the given context. Please provide a short and concise response."
+    }
+    
+    response_generation_params = {
+        "temperature": 0,
+        "top_p": 0.9,
+        "stop": ["[EOS]", "</s>"],
+        "skip_special_tokens": False,
+        "include_stop_str_in_output": True,
+        "logprobs": 1,
+        "spaces_between_special_tokens": False,
+        "max_tokens": 4096
+    }
+    
+    other_generation_params = {
+        "temperature": 1,
+        "top_p": 0.9,
+        "stop": ["[EOS]", "</s>"],
+        "skip_special_tokens": False,
+        "include_stop_str_in_output": True,
+        "logprobs": 1,
+        "spaces_between_special_tokens": False,
+        "max_tokens": 4096
+    }
+
+    def __init__(
+        self,
+        config: dict,
+        prompt_template = None,
+        retriever = None,
+        generator = None,
+        max_depth = 3,
+        batch_size = 32
+    ):
+        super().__init__(config, prompt_template)
+
+
+        self.generator = generator if generator is not None else get_generator(config)
+        self.tokenizer = AutoTokenizer.from_pretrained(config["generator_model_path"], padding_side = "left")
+        self.retriever = retriever if retriever is not None else get_retriever(config)
+        
+        self.max_depth = max_depth
+        self.batch_size = batch_size
+        
+        # Due to the low effiency of original method, it only supports vllm now.
+    
+    def preprocess_eval_data(self, items: List) -> List[str]:
+        eval_examples = []
+
+        for item in items:
+            eval_example = f"<s><|system|>\n{self.system_prompt['qa']}" + self.tokenizer.eos_token + "\n<|user|>\n" + item.question + self.tokenizer.eos_token + "\n"
+            eval_example += "<|assistant|>\n"
+            eval_examples.append(eval_example)
+
+        return eval_examples
+
+    def format_evidences(self, evidences: List[str]):
+        format_evidence = ""
+        for evidence in evidences:
+            title = evidence['contents'].split('\n')[0]
+            text = "\n".join(evidence['contents'].split('\n')[1:])
+            format_evidence += f"Title: {title}\n"
+            format_evidence += f"Text: {text}\n"
+        return format_evidence
+
+    def generate_tree_of_thoughts_batch(self, initial_prompts_batch: List[str]):
+        paths_batch_dict = {
+            idx: [{
+                "prompt": initial_prompt,
+                "depth": 0,
+                "done": False
+            }]
+            for idx, initial_prompt in enumerate(initial_prompts_batch)
+        }
+        
+        final_outputs_batch = {idx: [] for idx in range(len(initial_prompts_batch))}
+        
+        while any(paths for paths in paths_batch_dict.values()):
+            current_batch = []
+            for i, _ in paths_batch_dict.items():
+                if paths_batch_dict[i]:
+                    current_path = paths_batch_dict[i].pop(0)
+                    current_batch.append(current_path)
+                else:
+                    continue
+            
+            if not current_batch:
+                break
+            
+            for special_token in self.expand_on_tokens:
+                
+                if current_batch[0]["depth"] >= self.max_depth and special_token != "[A_Response]":
+                    continue
+                
+                # Prepare for inputs
+                input_texts = [path["prompt"] + special_token for path in current_batch]
+            
+                # Generate outputs
+                if special_token != "[A_Response]":
+                    init_outputs = self.generator.generate(
+                        input_list = input_texts,
+                        return_raw_output = True,
+                        **self.response_generation_params
+                    )
+                else:
+                    init_outputs = self.generator.generate(
+                        input_list = input_texts,
+                        return_raw_output = True,
+                        **self.other_generation_params
+                    )
+
+                # Decode outputs
+                decoded_outputs = [output.outputs[0].text for output in init_outputs]
+                # Initialize lists to collect queries for batch retrieval
+                queries_for_search = []
+                
+                # Process outputs and prepare for retrieval
+                for i, decoded_output in enumerate(decoded_outputs):
+                    current_path = current_batch[i]
+                    decoded_output = decoded_output.replace("<s> ", "<s>")
+                    
+                    if special_token == "[A_Response]":
+                        pattern = r"(.*?)\[EOS\]"
+                        matches = re.findall(pattern, decoded_output, re.DOTALL)
+                        result = matches[-1].strip() if matches else "Unable to detect valid answer"
+                        token_ids = init_outputs[i].outputs[0].token_ids[1:-1]
+                        logprobs = init_outputs[i].outputs[0].logprobs[1:-1]
+                        confidence = 0
+                        for token_id, logprobs in zip(token_ids, logprobs):
+                            logprob = logprobs[token_id].logprob
+                            prob = math.exp(logprob)
+                            confidence += prob
+                        
+                        if len(token_ids) > 0:
+                            confidence /= len(token_ids)
+                        
+                        new_path = {
+                            "prompt": input_texts[i] + decoded_output,
+                            "depth": current_path["depth"] + 1,
+                            "done": True,
+                            "final_answer": result,
+                            "confidence": confidence
+                        }
+                        final_outputs_batch[i].append(new_path)
+                    else:
+                        # Extract the query
+                        pattern = r"(.*?)\[EOS\]"
+                        matches = re.findall(pattern, decoded_output, re.DOTALL)
+                        query_for_search = matches[-1].strip() if matches else "dummy"
+                        queries_for_search.append(query_for_search)
+                
+                # Perform batch retrieval
+                if queries_for_search:
+                    batch_search_results = self.retriever.batch_search(queries_for_search)
+                    
+                    for i, decoded_output in enumerate(decoded_outputs):
+                        search_results = batch_search_results[i]
+                        format_evidence = self.format_evidences(search_results)
+                        new_prompt = decoded_output + "[R_Evidences]" + format_evidence + "[/R_Evidences]"
+                        new_path = {
+                            "prompt": input_texts[i] + new_prompt,
+                            "depth": current_path["depth"] + 1,
+                            "done": False,
+                        }
+                        paths_batch_dict[i].append(new_path)
+
+        final_outputs_batch_list = [final_outputs_batch[i] for i in range(len(initial_prompts_batch))]
+        
+        return final_outputs_batch_list
+
+    def select_best_path_single_turn(self, final_outputs):
+        # After generating all paths, we can select the best answer
+        # Compute perplexity and confidence for each path
+        
+        scores = []
+        for path in final_outputs:
+            confidence = path["confidence"]
+            path["confidence"] = confidence
+            scores.append((path, confidence))
+
+        # Select the path with the highest confidence
+        best_path = max(scores, key = lambda x: x[1])[0]  # x[2] is confidence
+        pred = best_path["final_answer"]
+
+        return pred, best_path
+
+    def run(self, dataset, do_eval = True):
+        preds = []
+        meta_results = []
+
+        from tqdm import tqdm
+        for i in tqdm(range(0, len(dataset), self.batch_size)):
+            batch_items = dataset[i : i + self.batch_size]
+            eval_datas = self.preprocess_eval_data(batch_items)
+            paths_batch = self.generate_tree_of_thoughts_batch(initial_prompts_batch = eval_datas)
+            for paths in paths_batch:
+                pred, best_path = self.select_best_path_single_turn(paths)
+                preds.append(pred)
+                meta_results.append(best_path)
+
+
+        dataset.update_output("paths", meta_results)
+        dataset.update_output("pred", preds)
+
+        dataset = self.evaluate(dataset, do_eval = do_eval)
+        return dataset
+    

{flashrag-dev-0.1.4.dev20241203 → flashrag-dev-0.1.4.dev20241210}/flashrag/prompt/base_prompt.py

@@ -23,7 +23,7 @@ class PromptTemplate:
             self.is_chat = False
             if "chat" in model_name or "instruct" in model_name:
                 self.is_chat = True
-            self.tokenizer = AutoTokenizer.from_pretrained(self.generator_path, trust_remote_code=True)
+            self.tokenizer = None
         else:
             self.is_chat = True
             self.enable_chat = True
@@ -79,6 +79,8 @@ class PromptTemplate:
             return truncated_messages
 
         else:
+            if self.tokenizer is None:
+                self.tokenizer = AutoTokenizer.from_pretrained(self.generator_path, trust_remote_code=True)
             assert isinstance(prompt, str)
             tokenized_prompt = self.tokenizer(prompt, truncation=False, return_tensors="pt").input_ids[0]
 

flashrag-dev-0.1.4.dev20241210/flashrag/version.py

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

{flashrag-dev-0.1.4.dev20241203 → flashrag-dev-0.1.4.dev20241210/flashrag_dev.egg-info}/PKG-INFO

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

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