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

{flashrag-dev-0.1.3.dev20241119/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.dev20241119
+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.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.dev20241119 → 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.dev20241119 → 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.dev20241119 → 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.dev20241119 → 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.dev20241119
+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.dev20241119/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.dev20241119/flashrag/version.py

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