flashrag-dev 0.1.1__py3-none-any.whl

flashrag/pipeline/branching_pipeline.py

@@ -0,0 +1,250 @@
+import itertools
+from typing import List
+import re
+from tqdm import tqdm
+import numpy as np
+from transformers import LogitsProcessorList
+from flashrag.utils import get_retriever, get_generator
+from flashrag.pipeline import BasicPipeline
+from flashrag.prompt import PromptTemplate
+
+
+class REPLUGPipeline(BasicPipeline):
+    def __init__(self, config, prompt_template=None):
+        from flashrag.pipeline.replug_utils import load_replug_model
+
+        super().__init__(config, prompt_template)
+        # load specify model for REPLUG
+        model = load_replug_model(config["generator_model_path"])
+        self.generator = get_generator(config, model=model)
+
+        self.retriever = get_retriever(config)
+
+    def build_single_doc_prompt(self, question: str, doc_list: List[str]):
+        return [self.prompt_template.get_string(question=question, formatted_reference=doc) for doc in doc_list]
+
+    def format_reference(self, doc_item):
+        content = doc_item["contents"]
+        title = content.split("\n")[0]
+        text = "\n".join(content.split("\n")[1:])
+        return f"Document(Title: {title}): {text}"
+
+    def run(self, dataset, do_eval=True, pred_process_fun=None):
+        import torch
+        from flashrag.pipeline.replug_utils import REPLUGLogitsProcessor
+
+        input_query = dataset.question
+
+        retrieval_results, doc_scores = self.retriever.batch_search(input_query, return_score=True)
+        dataset.update_output("retrieval_result", retrieval_results)
+        dataset.update_output("doc_scores", doc_scores)
+
+        pred_answer_list = []
+        # each doc has a prompt
+        for item in tqdm(dataset, desc="Inference: "):
+            docs = [self.format_reference(doc_item) for doc_item in item.retrieval_result]
+            prompts = self.build_single_doc_prompt(question=item.question, doc_list=docs)
+
+            scores = torch.tensor(item.doc_scores, dtype=torch.float32).to(self.device)
+            output = self.generator.generate(
+                prompts, batch_size=len(docs), logits_processor=LogitsProcessorList([REPLUGLogitsProcessor(scores)])
+            )
+            # the output of the batch is same
+            output = output[0]
+            pred_answer_list.append(output)
+
+        dataset.update_output("pred", pred_answer_list)
+
+        dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
+
+        return dataset
+
+
+class SuRePipeline(BasicPipeline):
+    def __init__(self, config, prompt_template=None):
+        super().__init__(config, prompt_template)
+        self.config = config
+        self.generator = get_generator(config)
+        self.retriever = get_retriever(config)
+
+        self.load_prompts()
+
+    def load_prompts(self):
+        # prompt for candidates generation
+        P_CAN_INSTRUCT = (
+            "Below are {N} passages related to the question at the end. After reading"
+            "the passages, provide two correct candidates for the answer to the"
+            "question at the end. Each answer should be in the form: (a) xx, (b)"
+            "yy, and should not exceed 3 words for each candidate.\n\n"
+            "{reference}"
+            "Question: {question}\n"
+            "Answer:"
+        )
+
+        # prompt for candidate-conditioned summarization
+        P_SUM_INSTRUCT = (
+            "Reference:\n{reference}\n"
+            "Your job is to act as a professional writer. You need to write a"
+            "good-quality passage that can support the given prediction about the"
+            "question only based on the information in the provided supporting passages.\n"
+            "Now, let's start. After you write, please write [DONE] to indicate you"
+            "are done. Do not write a prefix (e.g., 'Response:') while writing a passage.\n"
+            "Question: {question}\n"
+            "Prediction: {pred}\n"
+            "Passage:"
+        )
+
+        # prompt for instance-wise validation
+        P_VAL_INSTRUCT = (
+            "Question: {question}\n"
+            "Prediction: {pred}\n"
+            "Passage: {summary}\n"
+            "Does the passage correctly support the prediction? Choices: [True,False].\n"
+            "Answer:"
+        )
+
+        # prompt for pair-wise ranking
+        P_RANK_INSTRUCT = (
+            "Question: Given the following passages, determine which one provides a"
+            "more informative answer to the subsequent question.\n"
+            "Passage 1: {summary1}\n"
+            "Passage 2: {summary2}\n"
+            "Target Question: {question}\n"
+            "Your Task:\n"
+            "Identify which passage (Passage 1 or Passage 2) is more relevant and"
+            "informative to answer the question at hand. Choices: [Passage 1,Passage 2].\n"
+            "Answer:"
+        )
+
+        self.P_CAN_TEMPLATE = PromptTemplate(self.config, "", P_CAN_INSTRUCT)
+        self.P_SUM_TEMPLATE = PromptTemplate(self.config, "", P_SUM_INSTRUCT)
+        self.P_VAL_TEMPLATE = PromptTemplate(self.config, "", P_VAL_INSTRUCT)
+        self.P_RANK_TEMPLATE = PromptTemplate(self.config, "", P_RANK_INSTRUCT)
+
+    @staticmethod
+    def format_ref(titles, texts):
+        formatted_ref = ""
+        idx = 1
+        for title, text in zip(titles, texts):
+            formatted_ref += f"Passage #{idx} Title: {title}\n"
+            formatted_ref += f"Passage #{idx} Text: {text}\n"
+            formatted_ref += "\n"
+            idx += 1
+        return formatted_ref
+
+    @staticmethod
+    def parse_candidates(model_response):
+        """Parse candidates from model response"""
+        model_response = model_response.strip("\n").strip()
+        # r'\([a-z]\) ([^,]+)'
+        candidates = re.findall("\((\w+)\)\s*([^()]+)", model_response)
+        candidates = [cand[1].split("\n")[0].strip() for cand in candidates]
+        # post-process
+        candidates = [cand.replace(",", "").strip() for cand in candidates]
+        return candidates
+
+    @staticmethod
+    def parse_validation(model_response):
+        """Parse model's validation result into score based on the paper formula"""
+        model_response = model_response.strip().lower()
+        if "true" in model_response:
+            return 1
+        else:
+            return 0
+
+    @staticmethod
+    def parse_ranking(model_response):
+        """Parse model's pair ranking result into score"""
+        model_response = model_response.strip().lower()
+        if "passage 1" in model_response:
+            score = 1
+        elif "passage 2" in model_response:
+            score = 0
+        else:
+            score = 0.5
+        return score
+
+    def run(self, dataset, do_eval=True, pred_process_fun=None):
+        input_query = dataset.question
+
+        retrieval_results, doc_scores = self.retriever.batch_search(input_query, return_score=True)
+        dataset.update_output("retrieval_result", retrieval_results)
+
+        pred_answer_list = []
+        for item in tqdm(dataset, desc="Pipeline runing: "):
+            retrieval_result = item.retrieval_result
+            doc_num = len(retrieval_result)
+            # format all docs
+            for doc_item in retrieval_result:
+                if "title" not in doc_item or "text" not in doc_item:
+                    doc_item["title"] = doc_item["contents"].split("\n")[0]
+                    doc_item["text"] = "\n".join(doc_item["contents"].split("\n")[1:])
+            formatted_ref = self.format_ref(
+                titles=[i["title"] for i in retrieval_result], texts=[i["text"] for i in retrieval_result]
+            )
+            # get candidates
+
+            input_prompt = self.P_CAN_TEMPLATE.get_string(
+                N=doc_num, formatted_reference=formatted_ref, question=item.question
+            )
+            output = self.generator.generate([input_prompt])[0]
+            candidates = self.parse_candidates(output)
+            item.update_output("candidates", candidates)
+
+            if len(candidates) == 0:
+                print("No valid predictions!")
+                pred = ""
+                pred_answer_list.append(pred)
+                continue
+
+            # get summarization for each candidate
+            input_prompts = [
+                self.P_SUM_TEMPLATE.get_string(question=item.question, pred=cand, formatted_reference=formatted_ref)
+                for cand in candidates
+            ]
+
+            all_summary = self.generator.generate(input_prompts)
+            item.update_output("all_summary", all_summary)
+
+            # instance-wise validation
+            input_prompts = [
+                self.P_VAL_TEMPLATE.get_string(question=item.question, pred=cand, summary=summary)
+                for cand, summary in zip(candidates, all_summary)
+            ]
+            val_results = self.generator.generate(input_prompts)
+            val_scores = [self.parse_validation(res) for res in val_results]
+            item.update_output("val_scores", val_scores)
+
+            # pair-wise ranking
+            summary_num = len(all_summary)
+            score_matrix = np.zeros((summary_num, summary_num))
+            iter_idxs = list(itertools.permutations(range(summary_num), 2))
+            input_prompts = [
+                self.P_RANK_TEMPLATE.get_string(
+                    question=item.question, summary1=all_summary[idx_tuple[0]], summary2=all_summary[idx_tuple[1]]
+                )
+                for idx_tuple in iter_idxs
+            ]
+            ranking_output = self.generator.generate(input_prompts)
+            ranking_scores = [self.parse_ranking(res) for res in ranking_output]
+            for idx_tuple, score in zip(iter_idxs, ranking_scores):
+                score_matrix[idx_tuple[0], idx_tuple[1]] = score
+            ranking_scores = score_matrix.sum(axis=1).squeeze().tolist()  # ranking score for each summary
+            item.update_output("ranking_scores", ranking_scores)
+
+            # combine two scores as the final score for each summary
+            if not isinstance(ranking_scores, list):
+                ranking_scores = [ranking_scores]
+            if not isinstance(val_scores, list):
+                val_scores = [val_scores]
+            total_scores = [x + y for x, y in zip(val_scores, ranking_scores)]
+
+            best_idx = np.argmax(total_scores)
+            pred = candidates[best_idx]
+            pred_answer_list.append(pred)
+
+        dataset.update_output("pred", pred_answer_list)
+
+        dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
+
+        return dataset

flashrag/pipeline/pipeline.py

@@ -0,0 +1,248 @@
+from flashrag.evaluator import Evaluator
+from flashrag.dataset.utils import split_dataset, merge_dataset
+from flashrag.utils import get_retriever, get_generator, get_refiner, get_judger
+from flashrag.prompt import PromptTemplate
+
+
+class BasicPipeline:
+    """Base object of all pipelines. A pipeline includes the overall process of RAG.
+    If you want to implement a pipeline, you should inherit this class.
+    """
+
+    def __init__(self, config, prompt_template=None):
+        self.config = config
+        self.device = config["device"]
+        self.retriever = None
+        self.evaluator = Evaluator(config)
+        self.save_retrieval_cache = config["save_retrieval_cache"]
+        if prompt_template is None:
+            prompt_template = PromptTemplate(config)
+        self.prompt_template = prompt_template
+
+    def run(self, dataset):
+        """The overall inference process of a RAG framework."""
+        pass
+
+    def evaluate(self, dataset, do_eval=True, pred_process_fun=None):
+        """The evaluation process after finishing overall generation"""
+
+        if pred_process_fun is not None:
+            raw_pred = dataset.pred
+            processed_pred = [pred_process_fun(pred) for pred in raw_pred]
+            dataset.update_output("raw_pred", raw_pred)
+            dataset.update_output("pred", processed_pred)
+
+        if do_eval:
+            # evaluate & save result
+            eval_result = self.evaluator.evaluate(dataset)
+            print(eval_result)
+
+        # save retrieval cache
+        if self.save_retrieval_cache:
+            self.retriever._save_cache()
+
+        return dataset
+
+
+class SequentialPipeline(BasicPipeline):
+    def __init__(self, config, prompt_template=None, retriever=None, generator=None):
+        """
+        inference stage:
+            query -> pre-retrieval -> retriever -> post-retrieval -> generator
+        """
+
+        super().__init__(config, prompt_template)
+        if generator is None:
+            self.generator = get_generator(config)
+        else:
+            self.generator = generator
+
+        if retriever is None:
+            self.retriever = get_retriever(config)
+        else:
+            self.retriever = retriever
+
+        # TODO: add rewriter module
+
+        self.use_fid = config["use_fid"]
+
+        if config["refiner_name"] is not None:
+            self.refiner = get_refiner(config, self.retriever, self.generator)
+        else:
+            self.refiner = None
+
+    def naive_run(self, dataset, do_eval=True, pred_process_fun=None):
+        # direct generation without RAG
+        input_prompts = [self.prompt_template.get_string(question=q) for q in dataset.question]
+        dataset.update_output("prompt", input_prompts)
+
+        pred_answer_list = self.generator.generate(input_prompts)
+        dataset.update_output("pred", pred_answer_list)
+
+        dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
+        return dataset
+
+    def run(self, dataset, do_eval=True, pred_process_fun=None):
+        input_query = dataset.question
+
+        retrieval_results = self.retriever.batch_search(input_query)
+        dataset.update_output("retrieval_result", retrieval_results)
+
+        if self.refiner:
+            input_prompt_flag = self.refiner.input_prompt_flag
+            if "llmlingua" in self.refiner.name and input_prompt_flag:
+                # input prompt
+                input_prompts = [
+                    self.prompt_template.get_string(question=q, retrieval_result=r)
+                    for q, r in zip(dataset.question, dataset.retrieval_result)
+                ]
+                dataset.update_output("prompt", input_prompts)
+                input_prompts = self.refiner.batch_run(dataset)
+            else:
+                # input retrieval docs
+                refine_results = self.refiner.batch_run(dataset)
+                dataset.update_output("refine_result", refine_results)
+                input_prompts = [
+                    self.prompt_template.get_string(question=q, formatted_reference=r)
+                    for q, r in zip(dataset.question, refine_results)
+                ]
+
+        else:
+            input_prompts = [
+                self.prompt_template.get_string(question=q, retrieval_result=r)
+                for q, r in zip(dataset.question, dataset.retrieval_result)
+            ]
+        dataset.update_output("prompt", input_prompts)
+
+        if self.use_fid:
+            print("Use FiD generation")
+            input_prompts = []
+            for item in dataset:
+                q = item.question
+                docs = item.retrieval_result
+                input_prompts.append([q + " " + doc for doc in docs])
+        # delete used refiner to release memory
+        if self.refiner:
+            del self.refiner
+        pred_answer_list = self.generator.generate(input_prompts)
+        dataset.update_output("pred", pred_answer_list)
+
+        dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
+
+        return dataset
+
+
+class ConditionalPipeline(BasicPipeline):
+    def __init__(self, config, prompt_template=None):
+        """
+        inference stage:
+            query -> judger -> sequential pipeline or naive generate
+        """
+
+        super().__init__(config, prompt_template)
+        self.judger = get_judger(config)
+        self.generator = get_generator(config)
+        self.retriever = get_retriever(config)
+
+        self.sequential_pipeline = SequentialPipeline(
+            config, prompt_template, retriever=self.retriever, generator=self.generator
+        )
+
+        self.zero_shot_templete = PromptTemplate(
+            config=config,
+            system_prompt="Answer the question based on your own knowledge. \
+                            Only give me the answer and do not output any other words.",
+            user_prompt="Question: {question}",
+        )
+
+    def run(self, dataset, do_eval=True, pred_process_fun=None):
+        # judge_result: list of bool element, representing whether to use retrieval
+        judge_result = self.judger.judge(dataset)
+        dataset.update_output("judge_result", judge_result)
+
+        # split dataset based on judge_result
+        dataset_split = split_dataset(dataset, judge_result)
+        pos_dataset, neg_dataset = dataset_split[True], dataset_split[False]
+
+        pos_dataset = self.sequential_pipeline.run(pos_dataset, do_eval=False)
+        self.sequential_pipeline.prompt_template = self.zero_shot_templete
+        neg_dataset = self.sequential_pipeline.naive_run(neg_dataset, do_eval=False)
+
+        # merge datasets into original format
+        dataset = merge_dataset(dataset_split, judge_result)
+
+        dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
+
+        return dataset
+
+
+class AdaptivePipeline(BasicPipeline):
+    def __init__(
+        self,
+        config,
+        norag_template=None,
+        single_hop_prompt_template=None,
+        multi_hop_prompt_template=None,
+    ):
+        super().__init__(config)
+        # load adaptive classifier as judger
+        self.judger = get_judger(config)
+
+        generator = get_generator(config)
+        retriever = get_retriever(config)
+        self.generator = generator
+        self.retriever = retriever
+
+        # Load three pipeline for three types of query: naive/single-hop/multi-hop
+        from flashrag.pipeline import IRCOTPipeline
+
+        if norag_template is None:
+            norag_templete = PromptTemplate(
+                config=config,
+                system_prompt="Answer the question based on your own knowledge. Only give me the answer and do not output any other words.",
+                user_prompt="Question: {question}",
+            )
+        self.norag_pipeline = SequentialPipeline(
+            config,
+            prompt_template=norag_templete,
+            retriever=retriever,
+            generator=generator,
+        )
+
+        self.single_hop_pipeline = SequentialPipeline(
+            config,
+            prompt_template=single_hop_prompt_template,
+            retriever=retriever,
+            generator=generator,
+        )
+
+        self.multi_hop_pipeline = IRCOTPipeline(
+            config,
+            prompt_template=multi_hop_prompt_template,
+            retriever=retriever,
+            generator=generator,
+        )
+
+    def run(self, dataset, do_eval=True, pred_process_fun=None):
+        # judge_result: choice result representing which pipeline to use(e.g. A, B, C)
+        judge_result = self.judger.judge(dataset)
+        dataset.update_output("judge_result", judge_result)
+
+        # split dataset based on judge_result
+        dataset_split = split_dataset(dataset, judge_result)
+        for symbol, symbol_dataset in dataset_split.items():
+            if symbol == "A":
+                symbol_dataset = self.norag_pipeline.naive_run(symbol_dataset, do_eval=False)
+            elif symbol == "B":
+                symbol_dataset = self.single_hop_pipeline.run(symbol_dataset, do_eval=False)
+            elif symbol == "C":
+                symbol_dataset = self.multi_hop_pipeline.run(symbol_dataset, do_eval=False)
+            else:
+                assert False, "Unknown symbol!"
+
+        # merge datasets into original format
+        dataset = merge_dataset(dataset_split, judge_result)
+
+        dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
+
+        return dataset