flashrag-dev 0.1.3.dev20241117__tar.gz → 0.1.3.dev20241120__tar.gz

{flashrag-dev-0.1.3.dev20241117/flashrag_dev.egg-info → flashrag-dev-0.1.3.dev20241120}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.3.dev20241117
+Version: 0.1.3.dev20241120
 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.3.dev20241117 → flashrag-dev-0.1.3.dev20241120}/flashrag/dataset/dataset.py

@@ -174,9 +174,12 @@ class Dataset:
         """Save the dataset into the original format."""
 
         save_data = [item.to_dict() for item in self.data]
-
+        def custom_serializer(obj):
+            if isinstance(obj, np.float32):  
+                return float(obj)           
+            raise TypeError(f"Type {type(obj)} not serializable")
         with open(save_path, "w", encoding="utf-8") as f:
-            json.dump(save_data, f, indent=4)
+            json.dump(save_data, f, indent=4, default=custom_serializer)
 
     def __str__(self) -> str:
         """Return a string representation of the dataset with a summary of items."""

{flashrag-dev-0.1.3.dev20241117 → flashrag-dev-0.1.3.dev20241120}/flashrag/dataset/utils.py

@@ -13,6 +13,8 @@ def convert_numpy(obj: Union[Dict, list, np.ndarray, np.generic]) -> Any:
         return obj.tolist()  # Convert numpy arrays to lists
     elif isinstance(obj, (np.integer, np.floating)):
         return obj.item()  # Convert numpy scalars to native Python scalars
+    elif isinstance(obj, np.float32):
+        return float(obj)
     else:
         return obj  # Return the object as-is if it's neither a dict, list, nor numpy type
 

flashrag-dev-0.1.3.dev20241120/flashrag/generator/fid.py

@@ -0,0 +1,206 @@
+# Source: FiD official repo: https://github.com/facebookresearch/FiD
+# This software is released under Creative Commons public licenses.
+
+import torch
+import torch.nn as nn
+import transformers
+import types
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss
+import numpy as np
+
+class FiDT5(transformers.T5ForConditionalGeneration):
+    def __init__(self, config):
+        super().__init__(config)
+        self.wrap_encoder()
+
+    def forward_(self, **kwargs):
+        if 'input_ids' in kwargs:
+            kwargs['input_ids'] = kwargs['input_ids'].view(kwargs['input_ids'].size(0), -1)
+        if 'attention_mask' in kwargs:
+            kwargs['attention_mask'] = kwargs['attention_mask'].view(kwargs['attention_mask'].size(0), -1)
+        
+        return super(FiDT5, self).forward(
+            **kwargs
+        )
+
+    # input ids : bs, n, seq_len -> bs, n*seq_len
+    # We need to resize as B x (N * L) instead of (B * N) x L here
+    # because the T5 forward method uses the input tensors to infer
+    # dimensions used in the decoder.
+    # EncoderWrapper resizes the inputs as (B * N) x L.
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        if input_ids != None:
+            # (bs, n, seq_len) -> (bs, n*seq_len) 
+            # inputs might have already be resized in the generate method
+            if input_ids.dim() == 3:
+                self.encoder.n_passages = input_ids.size(1)
+                input_ids = input_ids.view(input_ids.size(0), -1)
+        if attention_mask != None:
+            attention_mask = attention_mask.view(attention_mask.size(0), -1)
+        #print(input_ids.shape)
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **kwargs
+        )
+
+    def generate(self, input_ids, attention_mask, **kwargs):
+        # input ids - bs, n, seq_len -> bs, n*seq_len
+        self.encoder.n_passages = input_ids.size(1)
+        return super().generate(
+            input_ids=input_ids.view(input_ids.size(0), -1),
+            attention_mask=attention_mask.view(attention_mask.size(0), -1),
+            **kwargs,
+        )
+
+    def wrap_encoder(self, use_checkpoint=False):
+        """
+        Wrap T5 encoder to obtain a Fusion-in-Decoder model.
+        """
+        self.encoder = EncoderWrapper(self.encoder)
+    
+ 
+    def unwrap_encoder(self):
+        """
+        Unwrap Fusion-in-Decoder encoder, useful to load T5 weights.
+        """
+        self.encoder = self.encoder.encoder
+        block = []
+        for mod in self.encoder.block:
+            block.append(mod.module)
+        block = nn.ModuleList(block)
+        self.encoder.block = block
+    
+    def load_t5(self, state_dict):
+        self.unwrap_encoder()
+        self.load_state_dict(state_dict)
+        self.wrap_encoder() 
+
+    def set_checkpoint(self, use_checkpoint):
+        """
+        Enable or disable checkpointing in the encoder.
+        See https://pytorch.org/docs/stable/checkpoint.html
+        """
+        for mod in self.encoder.encoder.block:
+            mod.use_checkpoint = use_checkpoint
+
+    def tie_weights(self):
+        pass  
+
+class CheckpointWrapper(torch.nn.Module):
+    """
+    Wrapper replacing None outputs by empty tensors, which allows the use of
+    checkpointing.
+    """
+    def __init__(self, module, use_checkpoint=False):
+        super().__init__()
+        self.module = module
+        self.use_checkpoint = use_checkpoint
+
+    def forward(self, hidden_states, attention_mask, position_bias, **kwargs):
+        if self.use_checkpoint and self.training:
+            kwargs = {k: v for k, v in kwargs.items() if v is not None}
+            def custom_forward(*inputs):
+                output = self.module(*inputs, **kwargs)
+                empty = torch.tensor(
+                    [],
+                    dtype=torch.float,
+                    device=output[0].device,
+                    requires_grad=True)
+                output = tuple(x if x is not None else empty for x in output)
+                return output
+
+            output = torch.utils.checkpoint.checkpoint(
+                custom_forward,
+                hidden_states,
+                attention_mask,
+                position_bias
+            )
+            output = tuple(x if x.size() != 0 else None for x in output)
+        else:
+            output = self.module(hidden_states, attention_mask, position_bias, **kwargs)
+        return output
+
+def apply_checkpoint_wrapper(t5stack, use_checkpoint):
+    """
+    Wrap each block of the encoder to enable checkpointing.
+    """
+    block = []
+    for mod in t5stack.block:
+        wrapped_mod = CheckpointWrapper(mod, use_checkpoint)
+        block.append(wrapped_mod)
+    block = nn.ModuleList(block)
+    t5stack.block = block
+
+    
+class FiDBart(transformers.BartForConditionalGeneration):
+    def __init__(self, config):
+        super().__init__(config)
+        self.wrap_encoder()        
+    
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        
+        if input_ids != None:
+            # (bs, n, seq_len) -> (bs, n*seq_len) 
+            # inputs might have already be resized in the generate method
+            if input_ids.dim() == 3:
+                self.model.encoder.n_passages = input_ids.size(1)
+                input_ids = input_ids.view(input_ids.size(0), -1)
+
+        if attention_mask != None:
+            attention_mask = attention_mask.view(attention_mask.size(0), -1) 
+        
+        return super().forward(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
+
+    def generate(self, input_ids, attention_mask, **kwargs):
+        self.model.encoder.n_passages = input_ids.size(1)
+        return super().generate(
+            input_ids=input_ids.view(input_ids.size(0), -1),
+            attention_mask=attention_mask.view(attention_mask.size(0),-1),
+            **kwargs)
+
+    def wrap_encoder(self):
+        """
+        Wrap T5 encoder to obtain a Fusion-in-Decoder model.
+        """
+        self.model.encoder = EncoderWrapper(self.model.encoder)
+        
+    def unwrap_encoder(self):
+        """
+        Unwrap Fusion-in-Decoder encoder, useful to load bart weights.
+        """
+        self.model.encoder = self.model.encoder.encoder
+        block = []
+        for mod in self.model.encoder.layers:
+            block.append(mod)
+        block = nn.ModuleList(block)
+        self.model.encoder.layers = block
+        
+    def load_pretrained_model(self, state_dict):
+        self.unwrap_encoder()
+        self.load_state_dict(state_dict) 
+        self.wrap_encoder() 
+    def tie_weights(self):
+        pass  
+
+class EncoderWrapper(torch.nn.Module):
+    def __init__(self, encoder,use_checkpoint=False):
+        super().__init__()
+        self.encoder = encoder
+        
+        try:
+            self.main_input_name = encoder.main_input_name
+        except:
+            pass
+        apply_checkpoint_wrapper(self.encoder, use_checkpoint)
+    
+    def forward(self, input_ids=None, attention_mask=None,**kwargs):
+        bsz, total_length = input_ids.shape
+        passage_length = total_length // self.n_passages
+        # total_input
+        input_ids = input_ids.view(bsz*self.n_passages, passage_length)
+        attention_mask = attention_mask.view(bsz*self.n_passages, passage_length)
+        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask, **kwargs) 
+        outputs.last_hidden_state = outputs.last_hidden_state.view(bsz, self.n_passages*passage_length, -1)
+        return outputs 

{flashrag-dev-0.1.3.dev20241117 → flashrag-dev-0.1.3.dev20241120}/flashrag/generator/generator.py

@@ -49,10 +49,9 @@ class EncoderDecoderGenerator(BaseGenerator):
         self.fid = config["use_fid"]
         model_config = AutoConfig.from_pretrained(self.model_path)
         arch = model_config.architectures[0].lower()
-        if "t5" in arch:
+        if "t5" in arch or 'fusionindecoder' in arch:
             if self.fid:
                 from flashrag.generator.fid import FiDT5
-
                 self.model = FiDT5.from_pretrained(self.model_path)
             else:
                 self.model = T5ForConditionalGeneration.from_pretrained(self.model_path)
@@ -65,17 +64,18 @@ class EncoderDecoderGenerator(BaseGenerator):
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
 
     def encode_passages(self, batch_text_passages: List[List[str]]):
+        # need size: [batch_size, passage_num, passage_len]
         passage_ids, passage_masks = [], []
-        for k, text_passages in enumerate(batch_text_passages):
-            p = self.tokenizer.batch_encode_plus(
+        for text_passages in batch_text_passages:
+            p = self.tokenizer(
                 text_passages,
                 max_length=self.max_input_len,
-                pad_to_max_length=True,
-                return_tensors="pt",
+                padding='max_length',
                 truncation=True,
+                return_tensors='pt'
             )
-            passage_ids.append(p["input_ids"][None])
-            passage_masks.append(p["attention_mask"][None])
+            passage_ids.append(p['input_ids'][None])
+            passage_masks.append(p['attention_mask'][None])
 
         passage_ids = torch.cat(passage_ids, dim=0)
         passage_masks = torch.cat(passage_masks, dim=0)
@@ -128,8 +128,15 @@ class EncoderDecoderGenerator(BaseGenerator):
                 ).to(self.device)
 
             # TODO: multi-gpu inference
-            outputs = self.model.generate(**inputs, **generation_params)
-
+            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:
+                outputs = self.model.generate(**inputs, **generation_params)
             outputs = self.tokenizer.batch_decode(
                 outputs,
                 skip_special_tokens=True,

{flashrag-dev-0.1.3.dev20241117 → flashrag-dev-0.1.3.dev20241120}/flashrag/pipeline/active_pipeline.py

@@ -898,12 +898,12 @@ class SelfAskPipeline(BasicPipeline):
         dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
         return dataset
 
-
+    
 class IRCOTPipeline(BasicPipeline):
     IRCOT_INSTRUCTION = 'You serve as an intelligent assistant, adept at facilitating users through complex, multi-hop reasoning across multiple documents. This task is illustrated through demonstrations, each consisting of a document set paired with a relevant question and its multi-hop reasoning thoughts. Your task is to generate one thought for current step, DON\'T generate the whole thoughts at once! If you reach what you believe to be the final step, start with "So the answer is:".'
     IRCOT_EXAMPLE = "Wikipedia Title: Kurram Garhi\nKurram Garhi is a small village located near the city of Bannu, which is the part of Khyber Pakhtunkhwa province of Pakistan. Its population is approximately 35000. Barren hills are near this village. This village is on the border of Kurram Agency. Other nearby villages are Peppal, Surwangi and Amandi Kala.\n\nWikipedia Title: 2001–02 UEFA Champions League second group stage\nEight winners and eight runners- up from the first group stage were drawn into four groups of four teams, each containing two group winners and two runners- up. Teams from the same country or from the same first round group could not be drawn together. The top two teams in each group advanced to the quarter- finals.\n\nWikipedia Title: Satellite tournament\nA satellite tournament is either a minor tournament or event on a competitive sporting tour or one of a group of such tournaments that form a series played in the same country or region.\n\nWikipedia Title: Trojkrsti\nTrojkrsti is a village in Municipality of Prilep, Republic of Macedonia.\n\nWikipedia Title: Telephone numbers in Ascension Island\nCountry Code:+ 247< br> International Call Prefix: 00 Ascension Island does not share the same country code( +290) with the rest of St Helena.\n\nQuestion: Are both Kurram Garhi and Trojkrsti located in the same country?\nThought: Kurram Garhi is located in the country of Pakistan. Trojkrsti is located in the country of Republic of Macedonia. Thus, they are not in the same country. So the answer is: no.\n\n"
 
-    def __init__(self, config, prompt_template=None, retriever=None, generator=None, max_iter=2):
+    def __init__(self, config, prompt_template=None, retriever=None, generator=None, max_iter=5):
         # if not provide prompt template, use default template provided by IRCOT
         if prompt_template is None:
             prompt_template = PromptTemplate(
@@ -920,60 +920,97 @@ class IRCOTPipeline(BasicPipeline):
 
         self.max_iter = max_iter
 
-    def run_item(self, item):
-        question = item.question
-        retrieval_result, scores = self.retriever.search(question, return_score=True)
-        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}
-
-        thoughts = []
+    def run_batch(self, items):
+        # Initialize the necessary data structures
+        batch_thoughts = {item_id: [] for item_id in range(len(items))}
         iter_num = 0
+        batch_retrieval_results = []
+        doc2score_batch = []
+        id2doc_batch = []
+
+        # Initial retrieval for all items in the batch
+        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)
+            doc2score_batch.append(doc2score)
+            id2doc_batch.append(id2doc)
+
+        # Start the iterative process
+        active_item_ids = list(range(len(items)))  # Track items that need more iterations
         while iter_num < self.max_iter:
-            input_prompt = self.prompt_template.get_string(
-                question=question, retrieval_result=retrieval_result, previous_gen=" ".join(thoughts)
-            )
-            new_thought = self.generator.generate(input_prompt,stop=['.', '\n'])[0]
-            thoughts.append(new_thought)
-            iter_num += 1
-            if "So the answer is:" in new_thought:
-                item.update_output(
-                    f"intermediate_output_iter{iter_num}",
-                    {
-                        "input_prompt": input_prompt,
-                        "new_thought": new_thought,
-                    },
+            # Generate prompts and new thoughts for the active items
+            input_prompts = [
+                self.prompt_template.get_string(
+                    question=items[item_id].question,
+                    retrieval_result=batch_retrieval_results[item_id],
+                    previous_gen=' '.join(batch_thoughts[item_id])
                 )
-                break
+                for item_id in active_item_ids
+            ]
 
-            # retrieve new docs and merge
-            new_retrieval_result, new_scores = self.retriever.search(new_thought, return_score=True)
-            for doc_item, score in zip(new_retrieval_result, new_scores):
-                id2doc[doc_item["id"]] = doc_item
-                doc_id = doc_item["id"]
-                if doc_id in doc2score:
-                    doc2score[doc_id] = max(doc2score[doc_id], score)
+            # Batch generation for active items
+            new_thoughts_batch = self.generator.generate(input_prompts, stop=['.', '\n'])
+            
+            # Update thoughts and determine next active items
+            new_active_item_ids = []
+            for idx, item_id in enumerate(active_item_ids):
+                new_thought = new_thoughts_batch[idx]
+                batch_thoughts[item_id].append(new_thought)
+                
+                # Check for termination condition
+                # Store intermediate outputs
+                if "So the answer is:" in new_thought:
+                    items[item_id].update_output(
+                        f'intermediate_output_iter{iter_num}', 
+                        {
+                            'input_prompt': input_prompts[idx],
+                            'new_thought': new_thought,
+                        },
+                    )
                 else:
-                    doc2score[doc_id] = score
-            sorted_doc_score = sorted(doc2score.items(), key=lambda x: x[1], reverse=False)
-            sorted_doc_id = [t[0] for t in sorted_doc_score]
-            retrieval_result = [id2doc[id] for id in sorted_doc_id]
-
-            item.update_output(
-                f"intermediate_output_iter{iter_num}",
-                {
-                    "input_prompt": input_prompt,
-                    "new_thought": new_thought,
-                    "new_retreival_result": new_retrieval_result,
-                },
-            )
+                    new_active_item_ids.append(item_id)
+
+            # Update active item IDs for the next iteration
+            active_item_ids = new_active_item_ids
+
+            # Perform batch retrieval for new thoughts of active items
+            if active_item_ids:
+                print(len(active_item_ids))
+                print(active_item_ids)
+                new_thoughts_for_retrieval = [batch_thoughts[item_id][-1] for item_id in active_item_ids]
+                new_retrieval_results, new_scoress = self.retriever.batch_search(new_thoughts_for_retrieval, return_score=True)
+
+                for i, item_id in enumerate(active_item_ids):
+                    new_retrieval_result, new_scores = new_retrieval_results[i],new_scoress[i]
+                    
+                    # Update doc2score and id2doc for the current item
+                    for doc_item, score in zip(new_retrieval_result, new_scores):
+                        doc_id = doc_item['id']
+                        id2doc_batch[item_id][doc_id] = doc_item
+                        if doc_id in doc2score_batch[item_id]:
+                            doc2score_batch[item_id][doc_id] = max(doc2score_batch[item_id][doc_id], score)
+                        else:
+                            doc2score_batch[item_id][doc_id] = score
 
-        item.update_output("retrieval_result", retrieval_result)
-        item.update_output("pred", " ".join(thoughts))
-        return item
+                    # Sort and update retrieval results
+                    sorted_doc_score = sorted(doc2score_batch[item_id].items(), key=lambda x: x[1], reverse=False)
+                    sorted_doc_id = [t[0] for t in sorted_doc_score]
+                    batch_retrieval_results[item_id] = [id2doc_batch[item_id][id] for id in sorted_doc_id]
+
+            iter_num += 1
+
+        # Final update for each item in the batch
+        for item_id, item in enumerate(items):
+            item.update_output('retrieval_result', batch_retrieval_results[item_id])
+            item.update_output('pred', ' '.join(batch_thoughts[item_id]))
 
     def run(self, dataset, do_eval=True, pred_process_fun=ircot_pred_parse):
-        for item in tqdm(dataset, desc="Inference: "):
-            self.run_item(item)
+
+        self.run_batch(dataset)
 
         dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
         return dataset

{flashrag-dev-0.1.3.dev20241117 → flashrag-dev-0.1.3.dev20241120}/flashrag/pipeline/pipeline.py

@@ -107,11 +107,11 @@ class SequentialPipeline(BasicPipeline):
                 ]
 
         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 not self.use_fid:
+                input_prompts = [
+                    self.prompt_template.get_string(question=q, retrieval_result=r)
+                    for q, r in zip(dataset.question, dataset.retrieval_result)
+                ]
 
         if self.use_fid:
             print("Use FiD generation")
@@ -119,7 +119,9 @@ class SequentialPipeline(BasicPipeline):
             for item in dataset:
                 q = item.question
                 docs = item.retrieval_result
-                input_prompts.append([q + " " + doc for doc in docs])
+                input_prompts.append([q + " " + doc['contents'] for doc in docs])
+        dataset.update_output("prompt", input_prompts)
+
         # delete used refiner to release memory
         if self.refiner:
             del self.refiner

{flashrag-dev-0.1.3.dev20241117 → flashrag-dev-0.1.3.dev20241120}/flashrag/retriever/retriever.py

@@ -7,9 +7,9 @@ from typing import List, Dict
 import functools
 from tqdm import tqdm
 import faiss
-
+import numpy as np
 from flashrag.utils import get_reranker
-from flashrag.retriever.utils import load_corpus, load_docs
+from flashrag.retriever.utils import load_corpus, load_docs, convert_numpy
 from flashrag.retriever.encoder import Encoder, STEncoder
 
 
@@ -136,8 +136,13 @@ class BaseRetriever:
                 self.cache = json.load(f)
 
     def _save_cache(self):
+        self.cache = convert_numpy(self.cache)
+        def custom_serializer(obj):
+            if isinstance(obj, np.float32):  
+                return float(obj)           
+            raise TypeError(f"Type {type(obj)} not serializable")
         with open(self.cache_save_path, "w") as f:
-            json.dump(self.cache, f, indent=4)
+            json.dump(self.cache, f, indent=4, default=custom_serializer)
 
     def _search(self, query: str, num: int, return_score: bool) -> List[Dict[str, str]]:
         r"""Retrieve topk relevant documents in corpus.

{flashrag-dev-0.1.3.dev20241117 → flashrag-dev-0.1.3.dev20241120}/flashrag/retriever/utils.py

@@ -1,8 +1,25 @@
 import json
 import warnings
+from typing import Dict, Any, Union, List, Dict
+import numpy as np
 import datasets
 from transformers import AutoTokenizer, AutoModel, AutoConfig
 
+def convert_numpy(obj: Union[Dict, list, np.ndarray, np.generic]) -> Any:
+    """Recursively convert numpy objects in nested dictionaries or lists to native Python types."""
+    if isinstance(obj, dict):
+        return {k: convert_numpy(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [convert_numpy(i) for i in obj]
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()  # Convert numpy arrays to lists
+    elif isinstance(obj, (np.integer, np.floating)):
+        return obj.item()  # Convert numpy scalars to native Python scalars
+    elif isinstance(obj, np.float32):
+        return float(obj)
+    else:
+        return obj  # Return the object as-is if it's neither a dict, list, nor numpy type
+
 
 def load_model(model_path: str, use_fp16: bool = False):
     model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)

{flashrag-dev-0.1.3.dev20241117 → flashrag-dev-0.1.3.dev20241120}/flashrag/utils/utils.py

@@ -36,7 +36,7 @@ def get_generator(config, **params):
     elif config["framework"] == "hf":
         model_config = AutoConfig.from_pretrained(config["generator_model_path"])
         arch = model_config.architectures[0]
-        if "t5" in arch.lower() or "bart" in arch.lower():
+        if "t5" in arch.lower() or "bart" in arch.lower() or 'fusionindecoder' in arch.lower():
             return getattr(importlib.import_module("flashrag.generator"), "EncoderDecoderGenerator")(config, **params)
         else:
             return getattr(importlib.import_module("flashrag.generator"), "HFCausalLMGenerator")(config, **params)

flashrag-dev-0.1.3.dev20241120/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.3dev20241120"

{flashrag-dev-0.1.3.dev20241117 → flashrag-dev-0.1.3.dev20241120/flashrag_dev.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.3.dev20241117
+Version: 0.1.3.dev20241120
 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.3.dev20241117/flashrag/generator/fid.py

@@ -1,247 +0,0 @@
-# Source: FiD official repo: https://github.com/facebookresearch/FiD
-# This software is released under Creative Commons public licenses.
-
-import types
-import torch
-import transformers
-import torch.nn.functional as F
-from torch import nn
-
-
-class FiDT5(transformers.T5ForConditionalGeneration):
-    def __init__(self, config):
-        super().__init__(config)
-        self.wrap_encoder()
-
-    def forward_(self, **kwargs):
-        if "input_ids" in kwargs:
-            kwargs["input_ids"] = kwargs["input_ids"].view(kwargs["input_ids"].size(0), -1)
-        if "attention_mask" in kwargs:
-            kwargs["attention_mask"] = kwargs["attention_mask"].view(kwargs["attention_mask"].size(0), -1)
-
-        return super(FiDT5, self).forward(**kwargs)
-
-    # We need to resize as B x (N * L) instead of (B * N) x L here
-    # because the T5 forward method uses the input tensors to infer
-    # dimensions used in the decoder.
-    # EncoderWrapper resizes the inputs as (B * N) x L.
-    def forward(self, input_ids=None, attention_mask=None, **kwargs):
-        if input_ids != None:
-            # inputs might have already be resized in the generate method
-            if input_ids.dim() == 3:
-                self.encoder.n_passages = input_ids.size(1)
-            input_ids = input_ids.view(input_ids.size(0), -1)
-        if attention_mask != None:
-            attention_mask = attention_mask.view(attention_mask.size(0), -1)
-        return super().forward(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
-
-    # We need to resize the inputs here, as the generate method expect 2D tensors
-    def generate(self, input_ids, attention_mask, max_length):
-        self.encoder.n_passages = input_ids.size(1)
-        return super().generate(
-            input_ids=input_ids.view(input_ids.size(0), -1),
-            attention_mask=attention_mask.view(attention_mask.size(0), -1),
-            max_length=max_length,
-        )
-
-    def wrap_encoder(self, use_checkpoint=False):
-        """
-        Wrap T5 encoder to obtain a Fusion-in-Decoder model.
-        """
-        self.encoder = EncoderWrapper(self.encoder, use_checkpoint=use_checkpoint)
-
-    def unwrap_encoder(self):
-        """
-        Unwrap Fusion-in-Decoder encoder, useful to load T5 weights.
-        """
-        self.encoder = self.encoder.encoder
-        block = []
-        for mod in self.encoder.block:
-            block.append(mod.module)
-        block = nn.ModuleList(block)
-        self.encoder.block = block
-
-    def load_t5(self, state_dict):
-        self.unwrap_encoder()
-        self.load_state_dict(state_dict)
-        self.wrap_encoder()
-
-    def set_checkpoint(self, use_checkpoint):
-        """
-        Enable or disable checkpointing in the encoder.
-        See https://pytorch.org/docs/stable/checkpoint.html
-        """
-        for mod in self.encoder.encoder.block:
-            mod.use_checkpoint = use_checkpoint
-
-    def reset_score_storage(self):
-        """
-        Reset score storage, only used when cross-attention scores are saved
-        to train a retriever.
-        """
-        for mod in self.decoder.block:
-            mod.layer[1].EncDecAttention.score_storage = None
-
-    def get_crossattention_scores(self, context_mask):
-        """
-        Cross-attention scores are aggregated to obtain a single scalar per
-        passage. This scalar can be seen as a similarity score between the
-        question and the input passage. It is obtained by averaging the
-        cross-attention scores obtained on the first decoded token over heads,
-        layers, and tokens of the input passage.
-
-        More details in Distilling Knowledge from Reader to Retriever:
-        https://arxiv.org/abs/2012.04584.
-        """
-        scores = []
-        n_passages = context_mask.size(1)
-        for mod in self.decoder.block:
-            scores.append(mod.layer[1].EncDecAttention.score_storage)
-        scores = torch.cat(scores, dim=2)
-        bsz, n_heads, n_layers, _ = scores.size()
-        # batch_size, n_head, n_layers, n_passages, text_maxlength
-        scores = scores.view(bsz, n_heads, n_layers, n_passages, -1)
-        scores = scores.masked_fill(~context_mask[:, None, None], 0.0)
-        scores = scores.sum(dim=[1, 2, 4])
-        ntokens = context_mask.sum(dim=[2]) * n_layers * n_heads
-        scores = scores / ntokens
-        return scores
-
-    def overwrite_forward_crossattention(self):
-        """
-        Replace cross-attention forward function, only used to save
-        cross-attention scores.
-        """
-        for mod in self.decoder.block:
-            attn = mod.layer[1].EncDecAttention
-            attn.forward = types.MethodType(cross_attention_forward, attn)
-
-
-class EncoderWrapper(torch.nn.Module):
-    """
-    Encoder Wrapper for T5 Wrapper to obtain a Fusion-in-Decoder model.
-    """
-
-    def __init__(self, encoder, use_checkpoint=False):
-        super().__init__()
-
-        self.encoder = encoder
-        apply_checkpoint_wrapper(self.encoder, use_checkpoint)
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        **kwargs,
-    ):
-        # total_length = n_passages * passage_length
-        bsz, total_length = input_ids.shape
-        passage_length = total_length // self.n_passages
-        input_ids = input_ids.view(bsz * self.n_passages, passage_length)
-        attention_mask = attention_mask.view(bsz * self.n_passages, passage_length)
-        outputs = self.encoder(input_ids, attention_mask, **kwargs)
-        outputs = (outputs[0].view(bsz, self.n_passages * passage_length, -1),) + outputs[1:]
-        return outputs
-
-
-class CheckpointWrapper(torch.nn.Module):
-    """
-    Wrapper replacing None outputs by empty tensors, which allows the use of
-    checkpointing.
-    """
-
-    def __init__(self, module, use_checkpoint=False):
-        super().__init__()
-        self.module = module
-        self.use_checkpoint = use_checkpoint
-
-    def forward(self, hidden_states, attention_mask, position_bias, **kwargs):
-        if self.use_checkpoint and self.training:
-            kwargs = {k: v for k, v in kwargs.items() if v is not None}
-
-            def custom_forward(*inputs):
-                output = self.module(*inputs, **kwargs)
-                empty = torch.tensor([], dtype=torch.float, device=output[0].device, requires_grad=True)
-                output = tuple(x if x is not None else empty for x in output)
-                return output
-
-            output = torch.utils.checkpoint.checkpoint(custom_forward, hidden_states, attention_mask, position_bias)
-            output = tuple(x if x.size() != 0 else None for x in output)
-        else:
-            output = self.module(hidden_states, attention_mask, position_bias, **kwargs)
-        return output
-
-
-def apply_checkpoint_wrapper(t5stack, use_checkpoint):
-    """
-    Wrap each block of the encoder to enable checkpointing.
-    """
-    block = []
-    for mod in t5stack.block:
-        wrapped_mod = CheckpointWrapper(mod, use_checkpoint)
-        block.append(wrapped_mod)
-    block = nn.ModuleList(block)
-    t5stack.block = block
-
-
-def cross_attention_forward(
-    self,
-    input,
-    mask=None,
-    kv=None,
-    position_bias=None,
-    past_key_value_state=None,
-    head_mask=None,
-    query_length=None,
-    use_cache=False,
-    output_attentions=False,
-):
-    """
-    This only works for computing cross attention over the input
-    """
-    assert kv != None
-    assert head_mask == None
-    assert position_bias != None or self.has_relative_attention_bias
-
-    bsz, qlen, dim = input.size()
-    n_heads, d_heads = self.n_heads, self.d_kv
-    klen = kv.size(1)
-
-    q = self.q(input).view(bsz, -1, n_heads, d_heads).transpose(1, 2)
-    if past_key_value_state == None:
-        k = self.k(kv).view(bsz, -1, n_heads, d_heads).transpose(1, 2)
-        v = self.v(kv).view(bsz, -1, n_heads, d_heads).transpose(1, 2)
-    else:
-        k, v = past_key_value_state
-
-    scores = torch.einsum("bnqd,bnkd->bnqk", q, k)
-
-    if mask is not None:
-        scores += mask
-
-    if position_bias is None:
-        position_bias = self.compute_bias(qlen, klen)
-    scores += position_bias
-
-    if self.score_storage is None:
-        self.score_storage = scores
-
-    attn = F.softmax(scores.float(), dim=-1).type_as(scores)
-    attn = F.dropout(attn, p=self.dropout, training=self.training)
-
-    output = torch.matmul(attn, v)
-    output = output.transpose(1, 2).contiguous().view(bsz, -1, self.inner_dim)
-    output = self.o(output)
-
-    if use_cache:
-        output = (output,) + ((k, v),)
-    else:
-        output = (output,) + (None,)
-
-    if output_attentions:
-        output = output + (attn,)
-
-    if self.has_relative_attention_bias:
-        output = output + (position_bias,)
-
-    return output

flashrag-dev-0.1.3.dev20241117/flashrag/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.3dev20241117"