Skip to content

FiD Model

rankify.generator.models.fid_model

BaseRAGModel

Bases: ABC

Base RAG Model for Retrieval-Augmented Generation (RAG).

This is an abstract base class for implementing LLM endpoints in rankify. It defines the interface for generating responses and optional embedding generation.

Methods:

Name Description
generate

str, **kwargs) -> str: Abstract method to generate a response based on the given prompt.
embed

str, **kwargs) -> List[float]: Optional method to generate embeddings for the given text.
Notes
  • This class serves as a blueprint for RAG models like OpenAIModel and HuggingFaceModel.
  • The embed method is optional and can be implemented if needed.
  • This class needs to be extended to include new LLM endpoints in Rankify.
Source code in rankify/generator/models/base_rag_model.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class BaseRAGModel(ABC):
    """
    **Base RAG Model** for Retrieval-Augmented Generation (RAG).

    This is an abstract base class for implementing LLM endpoints in rankify. 
    It defines the interface for generating responses and optional embedding generation.

    Methods:
        generate(prompt: str, **kwargs) -> str:
            Abstract method to generate a response based on the given prompt.
        embed(text: str, **kwargs) -> List[float]:
            Optional method to generate embeddings for the given text.

    Notes:
        - This class serves as a blueprint for RAG models like `OpenAIModel` and `HuggingFaceModel`.
        - The `embed` method is optional and can be implemented if needed.
        - This class needs to be extended to include new LLM endpoints in Rankify.
    """

    @abstractmethod
    def generate(self, prompt: str, **kwargs) -> str:
        """Generate a response based on the given prompt."""
        pass

    def embed(self, text: str, **kwargs) -> List[float]:
        """Optional: Generate embeddings for the given text."""
        raise NotImplementedError("Embedding is not required for this implementation.")

generate(prompt, **kwargs) abstractmethod

Generate a response based on the given prompt.

Source code in rankify/generator/models/base_rag_model.py 
23
24
25
26
@abstractmethod
def generate(self, prompt: str, **kwargs) -> str:
    """Generate a response based on the given prompt."""
    pass

embed(text, **kwargs)

Optional: Generate embeddings for the given text.

Source code in rankify/generator/models/base_rag_model.py 
28
29
30
def embed(self, text: str, **kwargs) -> List[float]:
    """Optional: Generate embeddings for the given text."""
    raise NotImplementedError("Embedding is not required for this implementation.")

Dataset

Bases: Dataset

Source code in rankify/utils/generator/FiD/data.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
class Dataset(torch.utils.data.Dataset):
    def __init__(self,
                 data,
                 n_context=None,
                 question_prefix='question:',
                 title_prefix='title:',
                 passage_prefix='context:'):
        self.data = data
        self.n_context = n_context
        self.question_prefix = question_prefix
        self.title_prefix = title_prefix
        self.passage_prefix = passage_prefix
        self.sort_data()

    def __len__(self):
        return len(self.data)

    def get_target(self, example):
        if 'target' in example:
            target = example['target']
            return target + ' </s>'
        elif 'answers' in example:
            return random.choice(example['answers']) + ' </s>'
        else:
            return None

    def __getitem__(self, index):
        example = self.data[index]
        question = self.question_prefix + " " + example['question']
        target = self.get_target(example)

        if 'ctxs' in example and self.n_context is not None:
            f = self.title_prefix + " {} " + self.passage_prefix + " {}"
            contexts = example['ctxs'][:self.n_context]
            passages = [f.format(c['title'], c['text']) for c in contexts]
            scores = [float(c['score']) for c in contexts]
            scores = torch.tensor(scores)
            # TODO(egrave): do we want to keep this?
            if len(contexts) == 0:
                contexts = [question]
        else:
            passages, scores = None, None


        return {
            'index' : index,
            'question' : question,
            'target' : target,
            'passages' : passages,
            'scores' : scores
        }

    def sort_data(self):
        if self.n_context is None or not 'score' in self.data[0]['ctxs'][0]:
            return
        for ex in self.data:
            ex['ctxs'].sort(key=lambda x: float(x['score']), reverse=True)

    def get_example(self, index):
        return self.data[index]

Collator

Bases: object

Source code in rankify/utils/generator/FiD/data.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class Collator(object):
    def __init__(self, text_maxlength, tokenizer, answer_maxlength=20):
        self.tokenizer = tokenizer
        self.text_maxlength = text_maxlength
        self.answer_maxlength = answer_maxlength

    def __call__(self, batch):
        assert(batch[0]['target'] != None)
        index = torch.tensor([ex['index'] for ex in batch])
        target = [ex['target'] for ex in batch]
        target = self.tokenizer.batch_encode_plus(
            target,
            max_length=self.answer_maxlength if self.answer_maxlength > 0 else None,
            pad_to_max_length=True,
            return_tensors='pt',
            truncation=True if self.answer_maxlength > 0 else False,
        )
        target_ids = target["input_ids"]
        target_mask = target["attention_mask"].bool()
        target_ids = target_ids.masked_fill(~target_mask, -100)

        def append_question(example):
            if example['passages'] is None:
                return [example['question']]
            return [example['question'] + " " + t for t in example['passages']]
        text_passages = [append_question(example) for example in batch]
        passage_ids, passage_masks = encode_passages(text_passages,
                                                     self.tokenizer,
                                                     self.text_maxlength)

        return (index, target_ids, target_mask, passage_ids, passage_masks)

FiDT5

Bases: T5ForConditionalGeneration

Source code in rankify/utils/generator/FiD/model.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
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.)
        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)

wrap_encoder(use_checkpoint=False)

Wrap T5 encoder to obtain a Fusion-in-Decoder model.

Source code in rankify/utils/generator/FiD/model.py 
57
58
59
60
61
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)

unwrap_encoder()

Unwrap Fusion-in-Decoder encoder, useful to load T5 weights.

Source code in rankify/utils/generator/FiD/model.py 
63
64
65
66
67
68
69
70
71
72
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

set_checkpoint(use_checkpoint)

Enable or disable checkpointing in the encoder. See https://pytorch.org/docs/stable/checkpoint.html

Source code in rankify/utils/generator/FiD/model.py 
79
80
81
82
83
84
85
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

reset_score_storage()

Reset score storage, only used when cross-attention scores are saved to train a retriever.

Source code in rankify/utils/generator/FiD/model.py 
87
88
89
90
91
92
93
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

get_crossattention_scores(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.

Source code in rankify/utils/generator/FiD/model.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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.)
    scores = scores.sum(dim=[1, 2, 4])
    ntokens = context_mask.sum(dim=[2]) * n_layers * n_heads
    scores = scores/ntokens
    return scores

overwrite_forward_crossattention()

Replace cross-attention forward function, only used to save cross-attention scores.

Source code in rankify/utils/generator/FiD/model.py 
120
121
122
123
124
125
126
127
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)

ModelDownloader

Utility class for downloading and extracting model files.

Source code in rankify/utils/generator/download.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class ModelDownloader:
    """
    Utility class for downloading and extracting model files.
    """

    @staticmethod
    def download_and_extract(url, output_dir):
        """
        Downloads and extracts a model from a given URL.
        """
        os.makedirs(output_dir, exist_ok=True)
        tar_path = os.path.join(output_dir, "model.tar.gz")

        response = requests.get(url, stream=True)
        total_size = int(response.headers.get("content-length", 0))
        with open(tar_path, "wb") as file, tqdm(
            desc="Downloading Model", total=total_size, unit="B", unit_scale=True
        ) as pbar:
            for chunk in response.iter_content(chunk_size=1024):
                if chunk:
                    file.write(chunk)
                    pbar.update(len(chunk))

        with tarfile.open(tar_path, "r:gz") as tar:
            subdir = tar.getnames()[0]  # Get the top-level directory inside the tarball
            tar.extractall(output_dir)

        # Move contents up if they were extracted into a subdirectory
        extracted_path = os.path.join(output_dir, subdir)
        if os.path.exists(extracted_path) and os.path.isdir(extracted_path):
            for file in os.listdir(extracted_path):
                os.rename(os.path.join(extracted_path, file), os.path.join(output_dir, file))
            os.rmdir(extracted_path)  # Remove the now-empty subdirectory

download_and_extract(url, output_dir) staticmethod

Downloads and extracts a model from a given URL.

Source code in rankify/utils/generator/download.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
@staticmethod
def download_and_extract(url, output_dir):
    """
    Downloads and extracts a model from a given URL.
    """
    os.makedirs(output_dir, exist_ok=True)
    tar_path = os.path.join(output_dir, "model.tar.gz")

    response = requests.get(url, stream=True)
    total_size = int(response.headers.get("content-length", 0))
    with open(tar_path, "wb") as file, tqdm(
        desc="Downloading Model", total=total_size, unit="B", unit_scale=True
    ) as pbar:
        for chunk in response.iter_content(chunk_size=1024):
            if chunk:
                file.write(chunk)
                pbar.update(len(chunk))

    with tarfile.open(tar_path, "r:gz") as tar:
        subdir = tar.getnames()[0]  # Get the top-level directory inside the tarball
        tar.extractall(output_dir)

    # Move contents up if they were extracted into a subdirectory
    extracted_path = os.path.join(output_dir, subdir)
    if os.path.exists(extracted_path) and os.path.isdir(extracted_path):
        for file in os.listdir(extracted_path):
            os.rename(os.path.join(extracted_path, file), os.path.join(output_dir, file))
        os.rmdir(extracted_path)  # Remove the now-empty subdirectory

Document

Represents a document consisting of a question, answers, and contexts.

Attributes:

Name Type Description
question Question

The question associated with the document.
answers Answer

The answers to the question.
contexts list[Context]

A list of related contexts.
reorder_contexts list[Context] or None

A reordered list of contexts based on relevance.
Source code in rankify/dataset/dataset.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
class Document:
    """
    Represents a document consisting of a question, answers, and contexts.

    Attributes:
        question (Question): The question associated with the document.
        answers (Answer): The answers to the question.
        contexts (list[Context]): A list of related contexts.
        reorder_contexts (list[Context] or None): A reordered list of contexts based on relevance.
    """
    def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
        """
        Initializes a Document instance.

        Args:
            question (Question): The question associated with the document.
            answers (Answer): The answers to the question.
            contexts (list[Context], optional): A list of contexts related to the question.

        Example:
            ```python
            q = Question("What is the capital of France?")
            a = Answer(["Paris"])
            c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
            d = Document(question=q, answers=a, contexts=[c1, c2])
            print(d)
            ```
        """
        self.question: Question = question
        self.answers: Answer = answers
        self.contexts: List[Context] = contexts
        self.reorder_contexts: List[Context] = None
        self.id = str(id) 

    @classmethod
    def from_dict(cls, data: dict,n_docs:int=100) -> 'Document':
        """
        Creates a Document instance from a dictionary.

        Args:
            data (dict): A dictionary containing the question, answers, and contexts.
            n_docs (int, optional): The number of contexts to include. Defaults to 100.

        Returns:
            Document: A new Document instance.

        Example:
            ```python
            data = {
                "question": "What is the capital of France?",
                "answers": ["Paris"],
                "ctxs": [
                    {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
                    {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
                ]
            }
            d = Document.from_dict(data)
            print(d.question)
            ```
        """
        question = Question(data["question"])
        if "answers" in data:
            answers = Answer(data["answers"])
        else:
            answers =Answer('')

        if "query_id" in data:
            id = data["query_id"]
        else:
            id = None
        contexts = [Context(**ctx) for ctx in data["ctxs"][:n_docs]]
        return cls(question, answers, contexts, id=id)

    def to_dict(self) -> Dict[str, Optional[object]]:
        """
        Converts the document into a dictionary representation.

        Returns:
            dict: A dictionary containing the question, answers, and contexts.
        """
        return {
            "question": self.question.question,
            "answers": self.answers.answers,
            "contexts": [ctx.to_dict() for ctx in self.contexts]
        }
    def to_dict_reoreder(self) -> Dict[str,Optional[object]]:
        return {
            "question" : self.question.question,
            "answers" : self.answers.answers,
            "contexts" : [ctx.to_dict() for ctx in self.reorder_contexts]
        }
    def __str__(self) -> str:
        """
        Returns a string representation of the Document instance.

        Returns:
            str: The formatted document information.

        Example:
            ```python
            d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
            print(d)
            ```
        """
        contexts_str = "\n\n".join([str(ctx) for ctx in self.contexts])
        reorder_contexts_str= ''
        if self.reorder_contexts is not None:
            reorder_contexts_str = "\n\n".join([str(ctx) for ctx in self.reorder_contexts])
        return f"{self.question}\n\n{self.answers}\n\nContext: \n\n{contexts_str}\nReorder contexts: \n\n{reorder_contexts_str}"

__init__(question, answers, contexts=None, id=None)

Initializes a Document instance.

Parameters:

Name Type Description Default
question Question

The question associated with the document.

 required
answers Answer

The answers to the question.

 required
contexts list[Context]

A list of contexts related to the question.

 None
Example 
q = Question("What is the capital of France?")
a = Answer(["Paris"])
c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
d = Document(question=q, answers=a, contexts=[c1, c2])
print(d)
Source code in rankify/dataset/dataset.py 
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
    """
    Initializes a Document instance.

    Args:
        question (Question): The question associated with the document.
        answers (Answer): The answers to the question.
        contexts (list[Context], optional): A list of contexts related to the question.

    Example:
        ```python
        q = Question("What is the capital of France?")
        a = Answer(["Paris"])
        c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
        d = Document(question=q, answers=a, contexts=[c1, c2])
        print(d)
        ```
    """
    self.question: Question = question
    self.answers: Answer = answers
    self.contexts: List[Context] = contexts
    self.reorder_contexts: List[Context] = None
    self.id = str(id)

from_dict(data, n_docs=100) classmethod

Creates a Document instance from a dictionary.

Parameters:

Name Type Description Default
data dict

A dictionary containing the question, answers, and contexts.

 required
n_docs int

The number of contexts to include. Defaults to 100.

 100

Returns:

Name Type Description
Document Document

A new Document instance.
Example 
data = {
    "question": "What is the capital of France?",
    "answers": ["Paris"],
    "ctxs": [
        {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
        {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
    ]
}
d = Document.from_dict(data)
print(d.question)
Source code in rankify/dataset/dataset.py 
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
@classmethod
def from_dict(cls, data: dict,n_docs:int=100) -> 'Document':
    """
    Creates a Document instance from a dictionary.

    Args:
        data (dict): A dictionary containing the question, answers, and contexts.
        n_docs (int, optional): The number of contexts to include. Defaults to 100.

    Returns:
        Document: A new Document instance.

    Example:
        ```python
        data = {
            "question": "What is the capital of France?",
            "answers": ["Paris"],
            "ctxs": [
                {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
                {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
            ]
        }
        d = Document.from_dict(data)
        print(d.question)
        ```
    """
    question = Question(data["question"])
    if "answers" in data:
        answers = Answer(data["answers"])
    else:
        answers =Answer('')

    if "query_id" in data:
        id = data["query_id"]
    else:
        id = None
    contexts = [Context(**ctx) for ctx in data["ctxs"][:n_docs]]
    return cls(question, answers, contexts, id=id)

to_dict()

Converts the document into a dictionary representation.

Returns:

Name Type Description
dict Dict[str, Optional[object]]

A dictionary containing the question, answers, and contexts.
Source code in rankify/dataset/dataset.py 
269
270
271
272
273
274
275
276
277
278
279
280
def to_dict(self) -> Dict[str, Optional[object]]:
    """
    Converts the document into a dictionary representation.

    Returns:
        dict: A dictionary containing the question, answers, and contexts.
    """
    return {
        "question": self.question.question,
        "answers": self.answers.answers,
        "contexts": [ctx.to_dict() for ctx in self.contexts]
    }

__str__()

Returns a string representation of the Document instance.

Returns:

Name Type Description
str str

The formatted document information.
Example 
d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
print(d)
Source code in rankify/dataset/dataset.py 
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
def __str__(self) -> str:
    """
    Returns a string representation of the Document instance.

    Returns:
        str: The formatted document information.

    Example:
        ```python
        d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
        print(d)
        ```
    """
    contexts_str = "\n\n".join([str(ctx) for ctx in self.contexts])
    reorder_contexts_str= ''
    if self.reorder_contexts is not None:
        reorder_contexts_str = "\n\n".join([str(ctx) for ctx in self.reorder_contexts])
    return f"{self.question}\n\n{self.answers}\n\nContext: \n\n{contexts_str}\nReorder contexts: \n\n{reorder_contexts_str}"

FiDModel

Bases: BaseRAGModel

FiD (Fusion-in-Decoder) Generator for Open-Domain Question Answering.

Fusion-in-Decoder (FiD) is a retrieval-augmented generation (RAG) model that aggregates information from multiple retrieved passages to generate context-aware answers. Inside Rankifys architecture, this model is an exception, it subclasses BaseRAGModel but also includes the logic of the FiD RAG method, since this RAG technique relies on the full transformer architecture.

Attributes:

Name Type Description
device str

Device used for inference ('cuda' or 'cpu').
model_path str

Path to the downloaded and extracted model.
tokenizer T5Tokenizer

T5 tokenizer for text encoding.
model FiDT5

Pretrained FiD model for text generation.
max_length int

Maximum length of the generated output (default: 50 tokens).
References
  • Izacard & Grave Leveraging Passage Retrieval with Generative Models for Open-Domain QA
    Paper
See Also
  • BaseRAGModel: Parent class for FiDModel.
  • RAG-based QA Models: FiD falls under retrieval-augmented generation.
Example 
from rankify.dataset.dataset import Document, Question, Answer, Context
from rankify.generator.generator import Generator

# Define question and answer
question = Question("What is the capital of France?")
answers = Answer([""])
contexts = [
    Context(id=1, title="France", text="The capital of France is Paris.", score=0.9),
    Context(id=2, title="Germany", text="Berlin is the capital of Germany.", score=0.5)
]

# Construct document
doc = Document(question=question, answers=answers, contexts=contexts)

# Initialize Generator (e.g., Meta Llama)
generator = Generator(method="fid", model_name='nq_reader_base', backend="fid")

# Generate answer
generated_answers = generator.generate([doc])
print(generated_answers)  # Output: ["Paris"]
Notes
  • FiD combines multiple passages to generate better responses.
  • It integrates seamlessly with the Generator class.
  • Uses retrieval-augmented generation (RAG) techniques for QA.
Source code in rankify/generator/models/fid_model.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
class FiDModel(BaseRAGModel):
    """
    **FiD (Fusion-in-Decoder) Generator** for Open-Domain Question Answering.


    **Fusion-in-Decoder (FiD)** is a **retrieval-augmented generation (RAG) model** that 
    aggregates information from **multiple retrieved passages** to generate context-aware answers.
    Inside Rankifys architecture, this model is an exception, it subclasses `BaseRAGModel` but also includes the logic
    of the FiD RAG method, since this RAG technique relies on the full transformer architecture.

    Attributes:
        device (str): Device used for inference (`'cuda'` or `'cpu'`).
        model_path (str): Path to the downloaded and extracted model.
        tokenizer (transformers.T5Tokenizer): T5 tokenizer for text encoding.
        model (FiDT5): Pretrained FiD model for text generation.
        max_length (int): Maximum length of the generated output (default: 50 tokens).

    References:
        - **Izacard & Grave** *Leveraging Passage Retrieval with Generative Models for Open-Domain QA*  
          [Paper](https://arxiv.org/abs/2007.01282)

    See Also:
        - `BaseRAGModel`: Parent class for FiDModel.
        - `RAG-based QA Models`: FiD falls under retrieval-augmented generation.

    Example:
        ```python
        from rankify.dataset.dataset import Document, Question, Answer, Context
        from rankify.generator.generator import Generator

        # Define question and answer
        question = Question("What is the capital of France?")
        answers = Answer([""])
        contexts = [
            Context(id=1, title="France", text="The capital of France is Paris.", score=0.9),
            Context(id=2, title="Germany", text="Berlin is the capital of Germany.", score=0.5)
        ]

        # Construct document
        doc = Document(question=question, answers=answers, contexts=contexts)

        # Initialize Generator (e.g., Meta Llama)
        generator = Generator(method="fid", model_name='nq_reader_base', backend="fid")

        # Generate answer
        generated_answers = generator.generate([doc])
        print(generated_answers)  # Output: ["Paris"]
        ```

    Notes:
        - FiD **combines multiple passages** to generate better responses.
        - It integrates **seamlessly** with the `Generator` class.
        - Uses **retrieval-augmented generation (RAG)** techniques for QA.
    """


    MODEL_URLS = {
        "nq_reader_base": "https://dl.fbaipublicfiles.com/FiD/pretrained_models/nq_reader_base.tar.gz",
        "nq_reader_large": "https://dl.fbaipublicfiles.com/FiD/pretrained_models/nq_reader_large.tar.gz",
        "tqa_reader_base": "https://dl.fbaipublicfiles.com/FiD/pretrained_models/tqa_reader_base.tar.gz",
        "tqa_reader_large": "https://dl.fbaipublicfiles.com/FiD/pretrained_models/tqa_reader_large.tar.gz"
    }

    CACHE_DIR = os.environ.get("RERANKING_CACHE_DIR", "./cache")

    def __init__(self, method: str = "fid", model_name: str = "nq_reader_base", **kwargs):
        """
        Initializes the FiDModel for retrieval-augmented generation.

        Args:
            method (str): The generator type (default: `"fid"`).
            model_name (str): The specific FiD model to use (e.g., `"nq_reader_base"`).
            max_length (int): Maximum length of generated answers (default: 50).
            device (str, optional): Device for inference (`'cuda'` or `'cpu'`). If None, auto-detect.
            **kwargs: Additional parameters for model configuration.

        Sets up device, downloads and loads the model and tokenizer, and configures generation parameters.

        Example:
            ```python
            generator = FiDModel(method="fid", model_name="nq_reader_base", max_length=64, device="cuda")
            ```
        """
        self.method = method
        self.model_name = model_name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model_path = self._download_model()
        self.tokenizer, self.model = self._load_model()
        self.max_length = kwargs.get("max_length", 50)  # Default max length for answers

    def _download_model(self) -> str:
        """
        Downloads and extracts the FiD model if not already present.

        Returns:
            str: The directory path where the model is stored.

        Raises:
            ValueError: If an unknown FiD model is specified.

        Example:
            ```python
            model_path = generator._download_model()
            ```
        """
        model_url = self.MODEL_URLS.get(self.model_name)
        if not model_url:
            raise ValueError(f"Unknown FiD model: {self.model_name}")

        model_dir = os.path.join(self.CACHE_DIR, "generator", "fid", self.model_name)
        if not os.path.exists(model_dir):
            ModelDownloader.download_and_extract(model_url, model_dir)

        return model_dir

    def _load_model(self):
        """
        Loads the FiD model and tokenizer.

        Returns:
            tuple: A tuple containing the tokenizer and the FiD model.

        Example:
            ```python
            tokenizer, model = generator._load_model()
            ```
        """
        tokenizer = transformers.T5Tokenizer.from_pretrained('t5-base', return_dict=False)
        model = FiDT5.from_pretrained(self.model_path)
        model = model.to(self.device)
        model.eval()
        return tokenizer, model

    def _prepare_dataloader(self, documents: list[Document]):
        """
        Converts `Document` objects into a FiD-compatible dataset.

        Args:
            documents (List[Document]): A list of documents with queries and retrieved contexts.

        Returns:
            tuple: A tuple containing a DataLoader and Dataset.

        Example:
            ```python
            dataloader, dataset = generator._prepare_dataloader(documents)
            ```
        """
        examples = []
        for doc in documents:
            example = {
                "question": doc.question.question,
                "answers": doc.answers.answers,
                "ctxs": [{"title": ctx.title, "text": ctx.text, "score": ctx.score} for ctx in doc.contexts]
            }
            examples.append(example)

        eval_dataset = Dataset(examples, n_context=len(documents[0].contexts))
        eval_sampler = SequentialSampler(eval_dataset)
        collator_function = Collator(512, self.tokenizer)  # 512 is max text length

        eval_dataloader = DataLoader(
            eval_dataset,
            sampler=eval_sampler,
            batch_size=4,  # Modify as needed
            num_workers=4,
            collate_fn=collator_function
        )

        return eval_dataloader, eval_dataset

    def generate(self, documents: list[Document]) -> list[str]:
        """
        Generates answers for a list of documents.

        Args:
            documents (List[Document]): A list of documents with queries and retrieved contexts.

        Returns:
            List[str]: A list of generated answers.

        Example:
            ```python
            generator = FiDGenerator(method="fid", model_name="nq_reader_base")
            documents = [Document(question=Question("Who wrote Hamlet?"))]
            print(generator.generate(documents))  # ['William Shakespeare wrote Hamlet in the early 1600s.']
            ```
        """
        eval_dataloader, eval_dataset = self._prepare_dataloader(documents)

        results = []
        with torch.no_grad():
            for batch in eval_dataloader:
                (idx, _, _, context_ids, context_mask) = batch

                outputs = self.model.generate(
                    input_ids=context_ids.to(self.device),
                    attention_mask=context_mask.to(self.device),
                    max_length=self.max_length
                )

                for k, output in enumerate(outputs):
                    answer = self.tokenizer.decode(output, skip_special_tokens=True)
                    results.append(answer)

        return results

__init__(method='fid', model_name='nq_reader_base', **kwargs)

Initializes the FiDModel for retrieval-augmented generation.

Parameters:

Name Type Description Default
method str

The generator type (default: "fid").

 'fid'
model_name str

The specific FiD model to use (e.g., "nq_reader_base").

 'nq_reader_base'
max_length int

Maximum length of generated answers (default: 50).

 required
device str

Device for inference ('cuda' or 'cpu'). If None, auto-detect.

 required
**kwargs

Additional parameters for model configuration.

 {}

Sets up device, downloads and loads the model and tokenizer, and configures generation parameters.

Example 
generator = FiDModel(method="fid", model_name="nq_reader_base", max_length=64, device="cuda")
Source code in rankify/generator/models/fid_model.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def __init__(self, method: str = "fid", model_name: str = "nq_reader_base", **kwargs):
    """
    Initializes the FiDModel for retrieval-augmented generation.

    Args:
        method (str): The generator type (default: `"fid"`).
        model_name (str): The specific FiD model to use (e.g., `"nq_reader_base"`).
        max_length (int): Maximum length of generated answers (default: 50).
        device (str, optional): Device for inference (`'cuda'` or `'cpu'`). If None, auto-detect.
        **kwargs: Additional parameters for model configuration.

    Sets up device, downloads and loads the model and tokenizer, and configures generation parameters.

    Example:
        ```python
        generator = FiDModel(method="fid", model_name="nq_reader_base", max_length=64, device="cuda")
        ```
    """
    self.method = method
    self.model_name = model_name
    self.device = "cuda" if torch.cuda.is_available() else "cpu"
    self.model_path = self._download_model()
    self.tokenizer, self.model = self._load_model()
    self.max_length = kwargs.get("max_length", 50)  # Default max length for answers

generate(documents)

Generates answers for a list of documents.

Parameters:

Name Type Description Default
documents List[Document]

A list of documents with queries and retrieved contexts.

 required

Returns:

Type Description
list[str]

List[str]: A list of generated answers.
Example 
generator = FiDGenerator(method="fid", model_name="nq_reader_base")
documents = [Document(question=Question("Who wrote Hamlet?"))]
print(generator.generate(documents))  # ['William Shakespeare wrote Hamlet in the early 1600s.']
Source code in rankify/generator/models/fid_model.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def generate(self, documents: list[Document]) -> list[str]:
    """
    Generates answers for a list of documents.

    Args:
        documents (List[Document]): A list of documents with queries and retrieved contexts.

    Returns:
        List[str]: A list of generated answers.

    Example:
        ```python
        generator = FiDGenerator(method="fid", model_name="nq_reader_base")
        documents = [Document(question=Question("Who wrote Hamlet?"))]
        print(generator.generate(documents))  # ['William Shakespeare wrote Hamlet in the early 1600s.']
        ```
    """
    eval_dataloader, eval_dataset = self._prepare_dataloader(documents)

    results = []
    with torch.no_grad():
        for batch in eval_dataloader:
            (idx, _, _, context_ids, context_mask) = batch

            outputs = self.model.generate(
                input_ids=context_ids.to(self.device),
                attention_mask=context_mask.to(self.device),
                max_length=self.max_length
            )

            for k, output in enumerate(outputs):
                answer = self.tokenizer.decode(output, skip_special_tokens=True)
                results.append(answer)

    return results