flashrag-dev 0.1.1__py3-none-any.whl

flashrag/generator/openai_generator.py

@@ -0,0 +1,96 @@
+import os
+from typing import List
+from copy import deepcopy
+import warnings
+from tqdm import tqdm
+import numpy as np
+
+import asyncio
+from openai import AsyncOpenAI, AsyncAzureOpenAI
+import tiktoken
+
+
+class OpenaiGenerator:
+    """Class for api-based openai models"""
+
+    def __init__(self, config):
+        self.model_name = config["generator_model"]
+        self.batch_size = config["generator_batch_size"]
+        self.generation_params = config["generation_params"]
+
+        self.openai_setting = config["openai_setting"]
+        if self.openai_setting["api_key"] is None:
+            self.openai_setting["api_key"] = os.getenv("OPENAI_API_KEY")
+
+        if "api_type" in self.openai_setting and self.openai_setting["api_type"] == "azure":
+            del self.openai_setting["api_type"]
+            self.client = AsyncAzureOpenAI(**self.openai_setting)
+        else:
+            self.client = AsyncOpenAI(**self.openai_setting)
+
+        self.tokenizer = tiktoken.encoding_for_model(self.model_name)
+
+    async def get_response(self, input: List, **params):
+        response = await self.client.chat.completions.create(model=self.model_name, messages=input, **params)
+        return response.choices[0]
+
+    async def get_batch_response(self, input_list: List[List], batch_size, **params):
+        total_input = [self.get_response(input, **params) for input in input_list]
+        all_result = []
+        for idx in tqdm(range(0, len(input_list), batch_size), desc="Generation process: "):
+            batch_input = total_input[idx : idx + batch_size]
+            batch_result = await asyncio.gather(*batch_input)
+            all_result.extend(batch_result)
+
+        return all_result
+
+    def generate(self, input_list: List[List], batch_size=None, return_scores=False, **params) -> List[str]:
+        # deal with single input
+        if len(input_list) == 1:
+            input_list = [input_list]
+        if batch_size is None:
+            batch_size = self.batch_size
+
+        # deal with generation params
+        generation_params = deepcopy(self.generation_params)
+        generation_params.update(params)
+        if "do_sample" in generation_params:
+            generation_params.pop("do_sample")
+
+        max_tokens = params.pop("max_tokens", None) or params.pop("max_new_tokens", None)
+        if max_tokens is not None:
+            generation_params["max_tokens"] = max_tokens
+        else:
+            generation_params["max_tokens"] = generation_params.get(
+                "max_tokens", generation_params.pop("max_new_tokens", None)
+            )
+        generation_params.pop("max_new_tokens", None)
+
+        if return_scores:
+            if generation_params.get("logprobs") is not None:
+                generation_params["logprobs"] = True
+                warnings.warn("Set logprobs to True to get generation scores.")
+            else:
+                generation_params["logprobs"] = True
+
+        if generation_params.get("n") is not None:
+            generation_params["n"] = 1
+            warnings.warn("Set n to 1. It can minimize costs.")
+        else:
+            generation_params["n"] = 1
+
+        loop = asyncio.get_event_loop()
+        result = loop.run_until_complete(self.get_batch_response(input_list, batch_size, **generation_params))
+
+        # parse result into response text and logprob
+        scores = []
+        response_text = []
+        for res in result:
+            response_text.append(res.message.content)
+            if return_scores:
+                score = np.exp(list(map(lambda x: x.logprob, res.logprobs.content)))
+                scores.append(score)
+        if return_scores:
+            return response_text, scores
+        else:
+            return response_text

flashrag/generator/stop_word_criteria.py

@@ -0,0 +1,105 @@
+"""
+Created by Nestor Demeure.
+This software is released under the Apache License 2.0.
+"""
+
+from typing import List
+import torch
+from transformers import StoppingCriteria, AutoTokenizer
+
+
+class StopWordCriteria(StoppingCriteria):
+    """
+    A stopping criteria that halts the text generation process if any specified stop word is encountered.
+
+    Inspired by https://discuss.huggingface.co/t/implimentation-of-stopping-criteria-list/20040/9
+    And: https://github.com/outlines-dev/outlines/blob/main/outlines/generate/api.py
+    """
+
+    def __init__(self, tokenizer: AutoTokenizer, prompts: List[str], stop_words: List[str] = [], check_every: int = 1):
+        """
+        Initializes the StopWordCriteria with the necessary parameters for checking stop words during text generation.
+
+        Parameters:
+            tokenizer (AutoTokenizer): The tokenizer for encoding prompts and stop words.
+            prompts (List[str]): Initial prompts used for generation, needed to determine where generated text begins.
+            stop_words (List[str]): Words that trigger the stopping of generation when detected.
+            check_every (int): Frequency of checking for stop words in the token stream (a performance optimization, use 1 to cut it out).
+        """
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.input_sizes = [self.tokenizer.encode(prompt, return_tensors="pt").size(-1) for prompt in prompts]
+        self.stop_words = stop_words
+        self.max_stop_word_size = max(
+            (self.tokenizer.encode(word, return_tensors="pt").size(-1) for word in stop_words), default=0
+        )
+        self.check_every = check_every
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        """
+        Determines whether to stop generation based on the presence of stop words.
+
+        Stops if a stop word is found in *all* batch elements *and* the sequence length is a multiple of `check_every`.
+        Note: Delay in stopping may occur if `check_every > 1`.
+
+        Parameters:
+            input_ids (torch.LongTensor): Generated token IDs.
+            scores (torch.FloatTensor): Generation scores for each token. Not used here.
+
+        Returns:
+            bool: True to stop generation, False to continue.
+        """
+        batch_size, seq_len = input_ids.shape
+
+        # Skip check if no stop words are defined or it is not yet time to check
+        if (len(self.stop_words) == 0) or (seq_len % self.check_every != 0):
+            return False
+
+        for i in range(batch_size):
+            # Calculate starting index for new tokens
+            prompt_size = self.input_sizes[i]
+            max_new_tokens = (2 * self.max_stop_word_size) + self.check_every
+            latest_tokens = input_ids[i, prompt_size:][-max_new_tokens:]
+
+            # Check for stop words in the decoded text
+            if not any(
+                word in self.tokenizer.decode(latest_tokens, skip_special_tokens=True) for word in self.stop_words
+            ):
+                return False  # Continue generation if any batch item lacks stop words
+
+        return True  # Stop generation if all conditions are met
+
+    def extract_answers(self, input_ids: torch.LongTensor, strip_stopword: bool = True) -> List[str]:
+        """
+        Extracts generated answers by removing prompts and optionally stopping at the first stop word.
+
+        Parameters:
+            input_ids (torch.LongTensor): Generated token IDs.
+            strip_stopword (bool): Determines whether the stop word is removed from the output.
+
+        Returns:
+            List[str]: Extracted answers, with or without stop words.
+        """
+        batch_size, _ = input_ids.shape
+        result = []
+
+        for i in range(batch_size):
+            # Decode generated tokens to text, excluding the prompt
+            prompt_size = self.input_sizes[i]
+            answer_tokens = input_ids[i, prompt_size:]
+            answer_text = self.tokenizer.decode(answer_tokens, skip_special_tokens=True)
+
+            # Find the first occurrence of any stop word
+            lower_stop_index = len(answer_text)  # Default to end of text
+            for word in self.stop_words:
+                stop_index = answer_text.find(word)
+                if stop_index != -1:
+                    # Adjust stop index based on whether we're stripping the stop word
+                    stop_index += 0 if strip_stopword else len(word)
+                    lower_stop_index = min(stop_index, lower_stop_index)
+
+            # Cut the text at the first stop word found (if any)
+            answer_text = answer_text[:lower_stop_index]
+            result.append(answer_text)
+
+        return result

flashrag/judger/__init__.py

@@ -0,0 +1 @@
+from flashrag.judger.judger import *

flashrag/judger/judger.py

@@ -0,0 +1,184 @@
+from typing import cast, List
+import json
+from tqdm.auto import trange
+from collections import Counter
+import numpy as np
+import torch
+import faiss
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from flashrag.retriever.utils import load_model, pooling
+
+
+class BaseJudger:
+    """Base object of Judger, used for judging whether to retrieve"""
+
+    def __init__(self, config):
+        self.config = config
+        self.name = config["judger_name"] if "judger_name" in config else None
+        self.judger_config = config["judger_config"] if "judger_config" in config else {}
+        self.device = config["device"]
+
+    def run(self, item) -> str:
+        """Get judgement result.
+
+        Args:
+            item: dataset item, contains question, retrieval result...
+
+        Returns:
+            judgement: bool, whether to retreive
+        """
+        pass
+
+    def batch_run(self, dataset, batch_size=None) -> List[str]:
+        return [self.run(item) for item in dataset]
+
+
+class SKRJudger(BaseJudger):
+    """Implementation for SKR-knn
+    Paper link: https://aclanthology.org/2023.findings-emnlp.691.pdf
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model_path = self.judger_config["model_path"]
+        self.training_data_path = self.judger_config["training_data_path"]
+        self.encoder, self.tokenizer = load_model(model_path=self.model_path, use_fp16=False)
+        self.topk = self.judger_config["topk"] if "topk" in self.judger_config else 5
+        self.batch_size = self.judger_config["batch_size"] if "batch_size" in config else 64
+        self.max_length = self.judger_config["max_length"] if "max_length" in config else 128
+
+        with open(self.training_data_path, "r") as f:
+            self.training_data = json.load(f)
+        # count number of pos & neg samples in training data
+        self.training_data_counter = Counter([item["judgement"].strip() for item in self.training_data])
+        self.training_pos_num = self.training_data_counter["ir_better"]
+        self.training_neg_num = self.training_data_counter["ir_worse"]
+        self.training_data_num = sum(self.training_data_counter.values())
+
+        # encode training question into faiss
+        training_questions = [item["question"] for item in self.training_data]
+        all_embeddings = self.encode(training_questions)
+        faiss_index = faiss.index_factory(all_embeddings.shape[-1], "Flat", faiss.METRIC_L2)
+        faiss_index.add(all_embeddings)
+        self.faiss = faiss_index
+
+    @torch.inference_mode(mode=True)
+    def encode(self, contents: list):
+        inputs = self.tokenizer(
+            contents,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+            max_length=self.max_length,
+        ).to("cuda")
+        output = self.encoder(**inputs, return_dict=True)
+        embeddings = pooling(output.pooler_output, output.last_hidden_state, inputs["attention_mask"], "pooler")
+
+        embeddings = cast(torch.Tensor, embeddings)
+        embeddings = torch.nn.functional.normalize(embeddings, dim=-1).detach()
+
+        all_embeddings = embeddings.cpu().numpy()
+        # all_embeddings = np.concatenate(all_embeddings, axis=0)
+        all_embeddings = all_embeddings.astype(np.float32)
+
+        return all_embeddings
+
+    def judge(self, dataset):
+        questions = dataset.question
+
+        all_judgements = []
+        for start_idx in range(0, len(questions), self.batch_size):
+            batch_question = questions[start_idx : start_idx + self.batch_size]
+            batch_emb = self.encode(batch_question)
+            scores, batch_idxs = self.faiss.search(batch_emb, k=self.topk)
+
+            for idxs in batch_idxs:
+                topk_samples = [self.training_data[idx]["judgement"].strip() for idx in idxs]
+                topk_counter = Counter(topk_samples)
+
+                # count number of pos & neg samples in topk
+                ir_better_num = topk_counter["ir_better"]
+                ir_worse_num = topk_counter["ir_worse"]
+                topk_delta = ir_better_num - ir_worse_num
+
+                training_data_delta = self.training_pos_num - self.training_neg_num
+
+                # provide judgments based on the formula in the paper
+                if training_data_delta < 0:
+                    if topk_delta < 0 and topk_delta <= int(training_data_delta * self.topk / self.training_data_num):
+                        judgement = False
+                    else:
+                        judgement = True
+                else:
+                    if topk_delta > 0 and topk_delta >= int(training_data_delta * self.topk / self.training_data_num):
+                        judgement = True
+                    else:
+                        judgement = False
+
+                all_judgements.append(judgement)
+
+        return all_judgements
+
+
+class AdaptiveJudger(BaseJudger):
+    """Implementation for Adaptive-RAG
+    Paper link: https://aclanthology.org/2024.naacl-long.389.pdf
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model_path = self.judger_config["model_path"]
+        self.batch_size = self.judger_config.get("batch_size", 16)
+        self.max_length = self.judger_config.get("max_length", 512)
+
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+        self.model.eval()
+        self.model.cuda()
+
+    @torch.inference_mode(mode=True)
+    def judge(self, dataset):
+        questions = dataset.question
+        questions = [q.strip() for q in questions]
+
+        all_preds = []
+        for idx in trange(0, len(questions), self.batch_size, desc="Judger process: "):
+            batch_input = questions[idx : idx + self.batch_size]
+            batch_input = self.tokenizer(
+                batch_input,
+                truncation=True,
+                padding=True,
+                max_length=512,
+                return_tensors="pt",
+            ).to(self.model.device)
+
+            scores = self.model.generate(
+                **batch_input, return_dict_in_generate=True, output_scores=True, max_length=self.max_length
+            ).scores[0]
+
+            probs = (
+                torch.nn.functional.softmax(
+                    torch.stack(
+                        [
+                            scores[:, self.tokenizer("A").input_ids[0]],
+                            scores[:, self.tokenizer("B").input_ids[0]],
+                            scores[:, self.tokenizer("C").input_ids[0]],
+                        ]
+                    ),
+                    dim=0,
+                )
+                .detach()
+                .cpu()
+                .numpy()
+            )
+
+            preds_labels = np.argmax(probs, 0)
+            label_to_option = {
+                0: "A",
+                1: "B",
+                2: "C",
+            }
+            preds = [label_to_option[pred] for pred in preds_labels]
+            all_preds.extend(preds)
+
+        return all_preds

flashrag/pipeline/__init__.py

@@ -0,0 +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