Skip to content

RankT5

rankify.models.rankt5

BaseRanking

Bases: ABC

An abstract base class for implementing different ranking models.

This class defines the interface for all ranking models, ensuring that all subclasses implement the required methods.

Attributes:

Name Type Description
method str

The name of the ranking method.
model_name str

The name of the model being used for ranking.
api_key str

An optional API key for accessing remote models or services.
Source code in rankify/models/base.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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class BaseRanking(ABC):
    """
    An abstract base class for implementing different ranking models.

    This class defines the interface for all ranking models, ensuring that all subclasses implement the required methods.

    Attributes:
        method (str): The name of the ranking method.
        model_name (str): The name of the model being used for ranking.
        api_key (str, optional): An optional API key for accessing remote models or services.
    """

    @abstractmethod
    def __init__(self, method: str= None, model_name: str= None, api_key: str= None, **kwargs) ->None:
        """
        Initializes the base ranking model.

        Args:
            method (str, optional): The name of the ranking method. Defaults to None.
            model_name (str, optional): The name of the model being used for ranking. Defaults to None.
            api_key (str, optional): An optional API key for accessing remote models or services. Defaults to None.

        Example:
            ```python
            class MyRanking(BaseRanking):
                def __init__(self, method, model_name):
                    super().__init__(method, model_name)
            ```
        """
        pass

    @abstractmethod
    def rank(self, documents: list[Document] ):
        """
        Abstract method to rank a list of documents.

        Args:
            documents (list[Document]): A list of Document instances that need to be ranked.

        Raises:
            NotImplementedError: This method must be implemented by subclasses.

        Example:
            ```python
            class MyRanking(BaseRanking):
                def __init__(self, method, model_name):
                    super().__init__(method, model_name)

                def rank(self, documents):
                    # Ranking implementation here
                    pass
            ```
        """
        pass

__init__(method=None, model_name=None, api_key=None, **kwargs) abstractmethod

Initializes the base ranking model.

Parameters:

Name Type Description Default
method str

The name of the ranking method. Defaults to None.

 None
model_name str

The name of the model being used for ranking. Defaults to None.

 None
api_key str

An optional API key for accessing remote models or services. Defaults to None.

 None
Example 
class MyRanking(BaseRanking):
    def __init__(self, method, model_name):
        super().__init__(method, model_name)
Source code in rankify/models/base.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
@abstractmethod
def __init__(self, method: str= None, model_name: str= None, api_key: str= None, **kwargs) ->None:
    """
    Initializes the base ranking model.

    Args:
        method (str, optional): The name of the ranking method. Defaults to None.
        model_name (str, optional): The name of the model being used for ranking. Defaults to None.
        api_key (str, optional): An optional API key for accessing remote models or services. Defaults to None.

    Example:
        ```python
        class MyRanking(BaseRanking):
            def __init__(self, method, model_name):
                super().__init__(method, model_name)
        ```
    """
    pass

rank(documents) abstractmethod

Abstract method to rank a list of documents.

Parameters:

Name Type Description Default
documents list[Document]

A list of Document instances that need to be ranked.

 required

Raises:

Type Description
NotImplementedError

This method must be implemented by subclasses.
Example 
class MyRanking(BaseRanking):
    def __init__(self, method, model_name):
        super().__init__(method, model_name)

    def rank(self, documents):
        # Ranking implementation here
        pass
Source code in rankify/models/base.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
@abstractmethod
def rank(self, documents: list[Document] ):
    """
    Abstract method to rank a list of documents.

    Args:
        documents (list[Document]): A list of Document instances that need to be ranked.

    Raises:
        NotImplementedError: This method must be implemented by subclasses.

    Example:
        ```python
        class MyRanking(BaseRanking):
            def __init__(self, method, model_name):
                super().__init__(method, model_name)

            def rank(self, documents):
                # Ranking implementation here
                pass
        ```
    """
    pass

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}"

RankT5

Bases: BaseRanking

Implements MonoRankT5, a T5-based re-ranking method that fine-tunes T5 with ranking losses.

RankT5 supports both MonoT5 and RankT5 re-ranking models:

  • MonoT5 applies a binary classification loss to determine query-document relevance.
  • RankT5 uses ranking losses to improve ranking effectiveness.
References
  • Zhuang, H. et al. (2023): RankT5: Fine-Tuning T5 for Text Ranking with Ranking Losses. Paper

Attributes:

Name Type Description
method str

The ranking method ("monot5" or "rankt5").
model_name str

The name of the T5 model.
model T5ForConditionalGeneration

The T5 model for ranking.
tokenizer T5Tokenizer

The tokenizer for encoding input texts.
max_input_length int

Maximum sequence length for input encoding.
batch_size int

The batch size for processing documents.
Example 
from rankify.dataset.dataset import Document, Question, Context
from rankify.models.reranking import Reranking

# Define a query and contexts
question = Question("What is the impact of climate change?")
contexts = [
    Context(text="Climate change causes rising sea levels and extreme weather.", id=0),
    Context(text="The stock market fluctuates due to various economic factors.", id=1),
    Context(text="Global warming contributes to increased wildfires and heatwaves.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize MonoRankT5
model = Reranking(method='rankt5', model_name='rankt5-base')
model.rank([document])

# Print reordered contexts
print("Reordered Contexts:")
for context in document.reorder_contexts:
    print(context.text)
Source code in rankify/models/rankt5.py 
 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
class RankT5(BaseRanking):
    """
    Implements **MonoRankT5**, a **T5-based re-ranking method** that fine-tunes **T5** with ranking losses.


    RankT5 supports both **MonoT5** and **RankT5** re-ranking models:

    - **MonoT5** applies a **binary classification loss** to determine query-document relevance.
    - **RankT5** uses **ranking losses** to improve ranking effectiveness.

    References:
        - **Zhuang, H. et al. (2023)**: *RankT5: Fine-Tuning T5 for Text Ranking with Ranking Losses*.
          [Paper](https://arxiv.org/abs/2210.10634)

    Attributes:
        method (str): The ranking method (`"monot5"` or `"rankt5"`).
        model_name (str): The name of the **T5 model**.
        model (T5ForConditionalGeneration): The **T5 model** for ranking.
        tokenizer (T5Tokenizer): The tokenizer for encoding input texts.
        max_input_length (int): Maximum sequence length for input encoding.
        batch_size (int): The batch size for processing documents.

    Example:
        ```python
        from rankify.dataset.dataset import Document, Question, Context
        from rankify.models.reranking import Reranking

        # Define a query and contexts
        question = Question("What is the impact of climate change?")
        contexts = [
            Context(text="Climate change causes rising sea levels and extreme weather.", id=0),
            Context(text="The stock market fluctuates due to various economic factors.", id=1),
            Context(text="Global warming contributes to increased wildfires and heatwaves.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize MonoRankT5
        model = Reranking(method='rankt5', model_name='rankt5-base')
        model.rank([document])

        # Print reordered contexts
        print("Reordered Contexts:")
        for context in document.reorder_contexts:
            print(context.text)
        ```
    """
    def __init__(self, method: str = None, model_name: str = 'castorini/monot5-base-msmarco', api_key: str = None , **kwargs):
        """
        Initializes the **MonoRankT5** class.

        Args:
            method (str): The ranking method (`"monot5"` or `"rankt5"`).
            model_name (str): The name of the **T5 model**.
            api_key (str, optional): Not used, included for framework consistency.

        Raises:
            ValueError: If the `method` is not `"monot5"` or `"rankt5"`.
        """
        self.method = method
        self.model_name = model_name
        self.mode = method.lower()
        if self.mode not in ['monot5', 'rankt5']:
            raise ValueError("Mode must be either 'monot5' or 'rankt5'")
        self.model, self.tokenizer = self.load_model_and_tokenizer()
        self.max_input_length = 512  # Default maximum input length
        self.batch_size = 10         # Default batch size for inference

    def load_model_and_tokenizer(self):
        """
        Loads the pre-trained **T5 model** and tokenizer.

        Returns:
            Tuple[T5ForConditionalGeneration, T5Tokenizer]: The loaded model and tokenizer.
        """
        print("Loading T5 model and tokenizer...")
        model = T5ForConditionalGeneration.from_pretrained(self.model_name).to('cuda')
        model.eval()
        tokenizer = T5Tokenizer.from_pretrained(self.model_name)
        print("Model and tokenizer loaded successfully.")
        return model, tokenizer

    def run_inference(self, input_tensors):
        """
        Runs inference using the **T5 model**.

        Args:
            input_tensors (dict): The input tensors for model inference.

        Returns:
            dict: The model outputs.
        """
        with torch.no_grad():
            output = self.model.generate(
                **input_tensors, 
                max_length=2,
                return_dict_in_generate=True,
                output_scores=True
            )
        return output

    def rank(self, documents: List[Document])-> List[Document]:
        """
        Reranks the passages for each document using **T5-based re-ranking**.

        Args:
            documents (List[Document]): List of documents containing queries and passages.

        Returns:
            List[Document]: The documents with updated `reorder_contexts`.
        """
        for document in tqdm(documents, desc="Reranking Documents"):
            query = document.question.question
            passages = copy.deepcopy(document.contexts)

            # Prepare input texts based on the selected mode
            if self.mode == 'monot5':
                input_texts = [f"Query: {query} Document: {x.text} Relevant:" for x in passages]
            elif self.mode == 'rankt5':
                input_texts = [f"Query: {query} Document: {x.text}" for x in passages]

            # Group inputs into batches
            grouped_input_texts = list(self.group2chunks(input_texts, n=self.batch_size))

            scores_holder = []
            for batch_input_texts in grouped_input_texts:
                input_tensors = self.tokenizer(
                    batch_input_texts, 
                    return_tensors='pt',
                    padding='max_length', 
                    max_length=self.max_input_length, 
                    truncation=True
                ).to('cuda')

                outputs = self.run_inference(input_tensors)
                del input_tensors

                scores = torch.stack(outputs.scores)

                # Extract scores based on the mode
                if self.mode == 'monot5':
                    # Extract scores for 'yes' (true) and 'no' (false)
                    yesno_softmax_scores = torch.nn.functional.log_softmax(
                        scores[0][:, [1176, 6136]], dim=1
                    )[:, 0].tolist()  # true, false
                    scores_holder += yesno_softmax_scores
                elif self.mode == 'rankt5':
                    # Extract the score associated with <extra_id_10>
                    rankt5_scores = scores[0][:, 32089].tolist()  # <extra_id_10>
                    scores_holder += rankt5_scores

            # Rank passages based on scores
            all_scores_tensor = torch.tensor(scores_holder)
            rank = torch.argsort(all_scores_tensor, descending=True).tolist()

            # Prepare the reordered list of contexts
            #reranked_passages = [copy.deepcopy(passages[rank_id]) for rank_id in rank]
            reranked_passages = []
            for new_idx in rank:
                context = copy.deepcopy(passages[new_idx])
                context.score = scores_holder[new_idx]
                reranked_passages.append(context)
            # Assign the reranked passages to reorder_contexts
            document.reorder_contexts = reranked_passages
        return documents
    def group2chunks(self, lst, n=5):
        """
        Groups a list into chunks of size **n**.

        Args:
            lst (list): The list to be chunked.
            n (int, optional): The chunk size (default is `5`).

        Yields:
            list: A chunk of `n` elements.
        """
        for i in range(0, len(lst), n):
            yield lst[i:i + n]

__init__(method=None, model_name='castorini/monot5-base-msmarco', api_key=None, **kwargs)

Initializes the MonoRankT5 class.

Parameters:

Name Type Description Default
method str

The ranking method ("monot5" or "rankt5").

 None
model_name str

The name of the T5 model.

 'castorini/monot5-base-msmarco'
api_key str

Not used, included for framework consistency.

 None

Raises:

Type Description
ValueError

If the method is not "monot5" or "rankt5".
Source code in rankify/models/rankt5.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(self, method: str = None, model_name: str = 'castorini/monot5-base-msmarco', api_key: str = None , **kwargs):
    """
    Initializes the **MonoRankT5** class.

    Args:
        method (str): The ranking method (`"monot5"` or `"rankt5"`).
        model_name (str): The name of the **T5 model**.
        api_key (str, optional): Not used, included for framework consistency.

    Raises:
        ValueError: If the `method` is not `"monot5"` or `"rankt5"`.
    """
    self.method = method
    self.model_name = model_name
    self.mode = method.lower()
    if self.mode not in ['monot5', 'rankt5']:
        raise ValueError("Mode must be either 'monot5' or 'rankt5'")
    self.model, self.tokenizer = self.load_model_and_tokenizer()
    self.max_input_length = 512  # Default maximum input length
    self.batch_size = 10         # Default batch size for inference

load_model_and_tokenizer()

Loads the pre-trained T5 model and tokenizer.

Returns:

Type Description

Tuple[T5ForConditionalGeneration, T5Tokenizer]: The loaded model and tokenizer.
Source code in rankify/models/rankt5.py 
81
82
83
84
85
86
87
88
89
90
91
92
93
def load_model_and_tokenizer(self):
    """
    Loads the pre-trained **T5 model** and tokenizer.

    Returns:
        Tuple[T5ForConditionalGeneration, T5Tokenizer]: The loaded model and tokenizer.
    """
    print("Loading T5 model and tokenizer...")
    model = T5ForConditionalGeneration.from_pretrained(self.model_name).to('cuda')
    model.eval()
    tokenizer = T5Tokenizer.from_pretrained(self.model_name)
    print("Model and tokenizer loaded successfully.")
    return model, tokenizer

run_inference(input_tensors)

Runs inference using the T5 model.

Parameters:

Name Type Description Default
input_tensors dict

The input tensors for model inference.

 required

Returns:

Name Type Description
dict

The model outputs.
Source code in rankify/models/rankt5.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def run_inference(self, input_tensors):
    """
    Runs inference using the **T5 model**.

    Args:
        input_tensors (dict): The input tensors for model inference.

    Returns:
        dict: The model outputs.
    """
    with torch.no_grad():
        output = self.model.generate(
            **input_tensors, 
            max_length=2,
            return_dict_in_generate=True,
            output_scores=True
        )
    return output

rank(documents)

Reranks the passages for each document using T5-based re-ranking.

Parameters:

Name Type Description Default
documents List[Document]

List of documents containing queries and passages.

 required

Returns:

Type Description
List[Document]

List[Document]: The documents with updated reorder_contexts.
Source code in rankify/models/rankt5.py 
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
def rank(self, documents: List[Document])-> List[Document]:
    """
    Reranks the passages for each document using **T5-based re-ranking**.

    Args:
        documents (List[Document]): List of documents containing queries and passages.

    Returns:
        List[Document]: The documents with updated `reorder_contexts`.
    """
    for document in tqdm(documents, desc="Reranking Documents"):
        query = document.question.question
        passages = copy.deepcopy(document.contexts)

        # Prepare input texts based on the selected mode
        if self.mode == 'monot5':
            input_texts = [f"Query: {query} Document: {x.text} Relevant:" for x in passages]
        elif self.mode == 'rankt5':
            input_texts = [f"Query: {query} Document: {x.text}" for x in passages]

        # Group inputs into batches
        grouped_input_texts = list(self.group2chunks(input_texts, n=self.batch_size))

        scores_holder = []
        for batch_input_texts in grouped_input_texts:
            input_tensors = self.tokenizer(
                batch_input_texts, 
                return_tensors='pt',
                padding='max_length', 
                max_length=self.max_input_length, 
                truncation=True
            ).to('cuda')

            outputs = self.run_inference(input_tensors)
            del input_tensors

            scores = torch.stack(outputs.scores)

            # Extract scores based on the mode
            if self.mode == 'monot5':
                # Extract scores for 'yes' (true) and 'no' (false)
                yesno_softmax_scores = torch.nn.functional.log_softmax(
                    scores[0][:, [1176, 6136]], dim=1
                )[:, 0].tolist()  # true, false
                scores_holder += yesno_softmax_scores
            elif self.mode == 'rankt5':
                # Extract the score associated with <extra_id_10>
                rankt5_scores = scores[0][:, 32089].tolist()  # <extra_id_10>
                scores_holder += rankt5_scores

        # Rank passages based on scores
        all_scores_tensor = torch.tensor(scores_holder)
        rank = torch.argsort(all_scores_tensor, descending=True).tolist()

        # Prepare the reordered list of contexts
        #reranked_passages = [copy.deepcopy(passages[rank_id]) for rank_id in rank]
        reranked_passages = []
        for new_idx in rank:
            context = copy.deepcopy(passages[new_idx])
            context.score = scores_holder[new_idx]
            reranked_passages.append(context)
        # Assign the reranked passages to reorder_contexts
        document.reorder_contexts = reranked_passages
    return documents

group2chunks(lst, n=5)

Groups a list into chunks of size n.

Parameters:

Name Type Description Default
lst list

The list to be chunked.

 required
n int

The chunk size (default is 5).

 5

Yields:

Name Type Description
list

A chunk of n elements.
Source code in rankify/models/rankt5.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
def group2chunks(self, lst, n=5):
    """
    Groups a list into chunks of size **n**.

    Args:
        lst (list): The list to be chunked.
        n (int, optional): The chunk size (default is `5`).

    Yields:
        list: A chunk of `n` elements.
    """
    for i in range(0, len(lst), n):
        yield lst[i:i + n]