Skip to content

Dataset

rankify.dataset.dataset

DownloadManger

Source code in rankify/utils/dataset/download.py 
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
class DownloadManger:
    @staticmethod
    def download(retriever: str, dataset: str, force_download: bool = True) ->str:
        cache_dir = get_cache_dir()


        if retriever not in HF_PRE_DEFIND_DATASET:
            raise FileNotFoundError(f"Retriever {retriever} Not Supported yet. Please choose another retriever.\nCheck Dataset.available_dataset()")
        if dataset not in HF_PRE_DEFIND_DATASET[retriever]:
            raise FileNotFoundError(f"Dataset {dataset} Not Supported yet. Please choose another dataset.\nCheck Dataset.available_dataset()")

        filename = HF_PRE_DEFIND_DATASET[retriever][dataset]['filename']
        if '-' in dataset:
            dataset_name, dataset_split = dataset.split('-', 1)
        else:
            dataset_name = dataset
        urls = HF_PRE_DEFIND_DATASET[retriever][dataset]['url']
        path = os.path.join(cache_dir, 'dataset', retriever, dataset_name)
        file_path = os.path.join(path, filename)

        # If force_download is False and file already exists, skip downloading
        if not force_download and os.path.exists(file_path):
            print(f"File {file_path} already exists. Skipping download.")
            return file_path

        os.makedirs(path, exist_ok=True)

        for url in urls:
            response = requests.get(url, stream=True)
            if response.status_code == 200:
                total_size = int(response.headers.get('content-length', 0))
                with open(file_path, 'wb') as file, tqdm(
                    desc=f"Downloading {retriever} {dataset_name} {filename}",
                    unit='B',
                    unit_scale=True,
                    unit_divisor=1024,
                    total=total_size,
                ) as bar:
                    # Update progress bar while streaming chunks
                    for chunk in response.iter_content(chunk_size=1024):
                        if chunk:  # Filter out keep-alive chunks
                            file.write(chunk)
                            bar.update(len(chunk))

                return file_path
            else:
                raise Exception(f'Failed to download the file from {url}')

Question

Represents a question with automatic validation.

Attributes:

Name Type Description
question str

The text of the question.
Source code in rankify/dataset/dataset.py 
 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
60
61
62
63
class Question:
    """
    Represents a question with automatic validation.

    Attributes:
        question (str): The text of the question.
    """
    def __init__(self, question: str ) -> None:
        """
        Initializes a Question instance.

        Args:
            question (str): The text of the question.

        Example:
            ```python
            q = Question("What is the capital of France?")
            print(q)  # Output: Question: What is the capital of France?
            ```
        """
        self.question = self.check_question(question)
    @classmethod
    def check_question(cls,question) -> str:
        """
        Ensures the question ends with a question mark.

        Args:
            question (str): The text of the question.

        Returns:
            str: The question with a question mark at the end if it was missing.

        Example:
            ```python
            Question.check_question("What is the capital of France")
            # Output: 'What is the capital of France?'
            ```
        """
        cls.question = question if question.endswith("?") else question +"?"
        return cls.question

    def __str__(self) -> str:
        """
        Returns a string representation of the Question instance.

        Returns:
            str: The formatted question.

        Example:
            ```python
            q = Question("What is the capital of France?")
            str(q)  # Output: 'Question: What is the capital of France?'
            ```
        """
        return f"Question: {self.question}"

__init__(question)

Initializes a Question instance.

Parameters:

Name Type Description Default
question str

The text of the question.

 required
Example 
q = Question("What is the capital of France?")
print(q)  # Output: Question: What is the capital of France?
Source code in rankify/dataset/dataset.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def __init__(self, question: str ) -> None:
    """
    Initializes a Question instance.

    Args:
        question (str): The text of the question.

    Example:
        ```python
        q = Question("What is the capital of France?")
        print(q)  # Output: Question: What is the capital of France?
        ```
    """
    self.question = self.check_question(question)

check_question(question) classmethod

Ensures the question ends with a question mark.

Parameters:

Name Type Description Default
question str

The text of the question.

 required

Returns:

Name Type Description
str str

The question with a question mark at the end if it was missing.
Example 
Question.check_question("What is the capital of France")
# Output: 'What is the capital of France?'
Source code in rankify/dataset/dataset.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
@classmethod
def check_question(cls,question) -> str:
    """
    Ensures the question ends with a question mark.

    Args:
        question (str): The text of the question.

    Returns:
        str: The question with a question mark at the end if it was missing.

    Example:
        ```python
        Question.check_question("What is the capital of France")
        # Output: 'What is the capital of France?'
        ```
    """
    cls.question = question if question.endswith("?") else question +"?"
    return cls.question

__str__()

Returns a string representation of the Question instance.

Returns:

Name Type Description
str str

The formatted question.
Example 
q = Question("What is the capital of France?")
str(q)  # Output: 'Question: What is the capital of France?'
Source code in rankify/dataset/dataset.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def __str__(self) -> str:
    """
    Returns a string representation of the Question instance.

    Returns:
        str: The formatted question.

    Example:
        ```python
        q = Question("What is the capital of France?")
        str(q)  # Output: 'Question: What is the capital of France?'
        ```
    """
    return f"Question: {self.question}"

Answer

Represents answers to a question.

Attributes:

Name Type Description
answers list[str]

A list of possible answers.
Source code in rankify/dataset/dataset.py 
 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
class Answer:
    """
    Represents answers to a question.

    Attributes:
        answers (list[str]): A list of possible answers.
    """
    def __init__(self, answers:list=None) -> None:
        """
        Initializes an Answer instance.

        Args:
            answers (list[str] or str, optional): A list of possible answers. Defaults to None.

        Example:
            ```python
            a = Answer(["Paris", "Lyon"])
            print(a)  
            # Output:
            # Answer:
            # - Paris
            # - Lyon
            ```
        """
        if isinstance(answers, str):  # If it's a string, convert it to a list
            self.answers = [answers]
        if isinstance(answers, int):
             self.answers = [str(answers)]
        elif isinstance(answers, list):  # If it's a list, ensure all elements are strings
            self.answers = [str(answer) for answer in answers]
        else:  # If it's neither, initialize with an empty list
            self.answers = []

    def __str__(self) -> str:
        """
        Returns a string representation of the Answer instance.

        Returns:
            str: The formatted answers.

        Example:
            ```python
            a = Answer(["Paris", "Lyon"])
            str(a)  
            # Output: 
            # Answer: 
            # - Paris
            # - Lyon
            ```
        """
        return f"Answer: \n- "+ "\n- ".join(self.answers)

__init__(answers=None)

Initializes an Answer instance.

Parameters:

Name Type Description Default
answers list[str] or str

A list of possible answers. Defaults to None.

 None
Example 
a = Answer(["Paris", "Lyon"])
print(a)  
# Output:
# Answer:
# - Paris
# - Lyon
Source code in rankify/dataset/dataset.py 
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
def __init__(self, answers:list=None) -> None:
    """
    Initializes an Answer instance.

    Args:
        answers (list[str] or str, optional): A list of possible answers. Defaults to None.

    Example:
        ```python
        a = Answer(["Paris", "Lyon"])
        print(a)  
        # Output:
        # Answer:
        # - Paris
        # - Lyon
        ```
    """
    if isinstance(answers, str):  # If it's a string, convert it to a list
        self.answers = [answers]
    if isinstance(answers, int):
         self.answers = [str(answers)]
    elif isinstance(answers, list):  # If it's a list, ensure all elements are strings
        self.answers = [str(answer) for answer in answers]
    else:  # If it's neither, initialize with an empty list
        self.answers = []

__str__()

Returns a string representation of the Answer instance.

Returns:

Name Type Description
str str

The formatted answers.
Example 
a = Answer(["Paris", "Lyon"])
str(a)  
# Output: 
# Answer: 
# - Paris
# - Lyon
Source code in rankify/dataset/dataset.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def __str__(self) -> str:
    """
    Returns a string representation of the Answer instance.

    Returns:
        str: The formatted answers.

    Example:
        ```python
        a = Answer(["Paris", "Lyon"])
        str(a)  
        # Output: 
        # Answer: 
        # - Paris
        # - Lyon
        ```
    """
    return f"Answer: \n- "+ "\n- ".join(self.answers)

Context

Represents a context with metadata such as score and title.

Attributes:

Name Type Description
score float

The relevance score of the context.
has_answer bool

Whether the context contains an answer.
id int

The identifier of the context.
title str

The title of the context.
text str

The text of the context.
Source code in rankify/dataset/dataset.py 
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
class Context:
    """
    Represents a context with metadata such as score and title.

    Attributes:
        score (float, optional): The relevance score of the context.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.
    """
    def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
        """
        Initializes a Context instance.

        Args:
            score (float, optional): The relevance score.
            has_answer (bool, optional): Whether the context contains an answer.
            id (int, optional): The identifier of the context.
            title (str, optional): The title of the context.
            text (str, optional): The text of the context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c)
            ```
        """
        self.score: Optional[float] = score
        self.has_answer: Optional[bool] = has_answer
        self.id: Optional[str] = id
        self.title: Optional[str] = title
        self.text: Optional[str] = text

    def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

        """
        Converts the Context instance to a dictionary.

        Args:
            save_text (bool): Whether to include text in the output dictionary.

        Returns:
            dict: The context data.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c.to_dict())
            ```
        """
        context_dict = {
            "score": float(self.score) if self.score is not None else None,
            "has_answer": self.has_answer,
            "id": self.id,
            }

        # Include 'text' only if save_text is True
        if save_text:
            context_dict["text"] = self.text
            context_dict["title"] =  self.title

        return context_dict
    def __str__(self) -> str:
        """
        Returns a string representation of the Context instance.

        Returns:
            str: The formatted context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(str(c))
            ```
        """
        return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

__init__(score=None, has_answer=None, id=None, title=None, text=None)

Initializes a Context instance.

Parameters:

Name Type Description Default
score float

The relevance score.

 None
has_answer bool

Whether the context contains an answer.

 None
id int

The identifier of the context.

 None
title str

The title of the context.

 None
text str

The text of the context.

 None
Example 
c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c)
Source code in rankify/dataset/dataset.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
    """
    Initializes a Context instance.

    Args:
        score (float, optional): The relevance score.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c)
        ```
    """
    self.score: Optional[float] = score
    self.has_answer: Optional[bool] = has_answer
    self.id: Optional[str] = id
    self.title: Optional[str] = title
    self.text: Optional[str] = text

to_dict(save_text=False)

Converts the Context instance to a dictionary.

Parameters:

Name Type Description Default
save_text bool

Whether to include text in the output dictionary.

 False

Returns:

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

The context data.
Example 
c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c.to_dict())
Source code in rankify/dataset/dataset.py 
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
def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

    """
    Converts the Context instance to a dictionary.

    Args:
        save_text (bool): Whether to include text in the output dictionary.

    Returns:
        dict: The context data.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c.to_dict())
        ```
    """
    context_dict = {
        "score": float(self.score) if self.score is not None else None,
        "has_answer": self.has_answer,
        "id": self.id,
        }

    # Include 'text' only if save_text is True
    if save_text:
        context_dict["text"] = self.text
        context_dict["title"] =  self.title

    return context_dict

__str__()

Returns a string representation of the Context instance.

Returns:

Name Type Description
str str

The formatted context.
Example 
c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(str(c))
Source code in rankify/dataset/dataset.py 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def __str__(self) -> str:
    """
    Returns a string representation of the Context instance.

    Returns:
        str: The formatted context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(str(c))
        ```
    """
    return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

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

Dataset

Represents a dataset for information retrieval.

Attributes:

Name Type Description
retriever str

The name of the retriever used to obtain the dataset.
dataset_name str

The name of the dataset.
n_docs int

The number of documents to include.
Source code in rankify/dataset/dataset.py 
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
class Dataset:
    """
    Represents a dataset for information retrieval.

    Attributes:
        retriever (str): The name of the retriever used to obtain the dataset.
        dataset_name (str): The name of the dataset.
        n_docs (int): The number of documents to include.
    """
    PASSAGES_URL = "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/psgs_w100/psgs_w100.tsv?download=true"
    CACHE_DIR = os.environ.get("RERANKING_CACHE_DIR", "./cache")
    PASSAGES_FILE = os.path.join(CACHE_DIR, "psgs_w100.tsv")
    def __init__(self, retriever: str, dataset_name:str, n_docs:int = 1000) -> None:
        """
        Initializes a Dataset instance.

        Args:
            retriever (str): The name of the retriever used to obtain the dataset.
            dataset_name (str): The name of the dataset.
            n_docs (int, optional): The number of documents to include. Defaults to 1000.

        Example:
            ```python
            dataset = Dataset(retriever='bm25', dataset_name='example_dataset', n_docs=500)
            print(dataset.dataset_name)
            ```
        """
        self.dataset_name: str  = dataset_name
        self.retriever: str  = retriever
        self.n_docs: int = n_docs
        self.documents: Optional[List[Document]]  = None
        self.id_to_text_title: Optional[Dict] = None

    def _ensure_passages_downloaded(self) -> None:
        """
        Ensures that the passages TSV file is downloaded and cached.
        """
        if not os.path.exists(self.PASSAGES_FILE):
            print(f"Downloading passages from {self.PASSAGES_URL}...")
            os.makedirs(os.path.dirname(self.PASSAGES_FILE), exist_ok=True)
            response = requests.get(self.PASSAGES_URL, stream=True)
            with open(self.PASSAGES_FILE, "wb") as f:
                for chunk in tqdm(response.iter_content(chunk_size=1024), desc="Downloading Passages"):
                    f.write(chunk)
            print(f"Passages file downloaded and saved to {self.PASSAGES_FILE}")
        else:
            print(f"Passages file already exists at {self.PASSAGES_FILE}")

    def _load_passages_mapping(self) -> Dict:
        """
        Loads the passages TSV file and creates a mapping of IDs to text and title.

        Returns
        -------
        Dict
            A dictionary mapping passage IDs to their text and title.
        """
        if self.id_to_text_title is None:
            self.id_to_text_title = {}
            with open(self.PASSAGES_FILE, "r", encoding="utf-8") as tsv_file:
                reader = csv.DictReader(tsv_file, delimiter="\t")
                for row in reader:
                    ctx_id = int(row["id"])
                    self.id_to_text_title[ctx_id] = {
                        "text": row["text"],
                        "title": row["title"]
                    }
        return self.id_to_text_title
    def update_contexts_from_passages(self) -> None:
        """
        Updates the text and title fields of contexts in all documents using the passages TSV file.
        If the TSV file is not already cached, it will be downloaded.

        Raises:
            ValueError: If the dataset has not been loaded before calling this method.

        Example:
            ```python
            dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
            dataset.download()
            dataset.update_contexts_from_passages()
            ```
        """
        if not self.documents:
            raise ValueError("Dataset has not been loaded. Call `download()` to load the dataset.")

        # Ensure TSV file is downloaded
        self._ensure_passages_downloaded()

        # Load the mapping of IDs to text and title
        id_to_text_title = self._load_passages_mapping()

        # Update contexts in all documents
        for document in tqdm(self.documents, desc="Retreive Text and Title for documents", unit="doc"):
            for context in document.contexts:
                if context.id in id_to_text_title:
                    context_data = id_to_text_title[context.id]
                    context.text = context_data["text"]
                    context.title = context_data["title"]


    def download(self, force_download:bool =True)-> List[Document]:
        """
        Downloads the dataset and loads it into memory.

        Args:
            force_download (bool, optional): Whether to force downloading the dataset even if it already exists locally. Defaults to True.

        Returns:
            list[Document]: A list of Document instances loaded from the dataset.

        Example:
            ```python
            dataset = Dataset(retriever='bm25', dataset_name='example_dataset', n_docs=500)
            documents = dataset.download()
            print(len(documents))
            ```
        """
        filepath= DownloadManger.download(self.retriever,self.dataset_name, force_download =force_download)
        self.documents= self.load_dataset(filepath, self.n_docs)
        if 'beir' not in self.dataset_name and 'dl' not in self.dataset_name:
            self.update_contexts_from_passages()
        return self.documents
    @classmethod
    def load_dataset(cls, filepath:str, n_docs: int= 100) -> List[Document]:
        """
        Loads the dataset from a JSON file.

        Args:
            filepath (str): The path to the JSON file containing the dataset.
            n_docs (int, optional): The number of documents to load. Defaults to 100.

        Returns:
            list[Document]: A list of Document instances loaded from the JSON file.

        Example:
            ```python
            filepath = 'example_dataset.json'
            documents = Dataset.load_dataset(filepath, n_docs=50)
            print(len(documents))
            ```
        """
        with open(filepath , encoding='utf-8') as file:
            data = json.load(file)
        data = [Document.from_dict(d,n_docs) for d in data]


        return data
    @classmethod
    def load_dataset_qa(cls, filepath: str) -> List[Document]:
        """
        Loads a QA dataset from a JSON or JSONL file. The dataset contains only questions, optionally answers, and optionally IDs.

        Args:
            filepath (str): The path to the JSON or JSONL file containing the dataset.

        Returns:
            List[Document]: A list of Document objects, each containing:
                - `question` (Question): The question object.
                - `answers` (Answer, optional): The answer object if available.
                - `id` (int, optional): The identifier if available.
                - `contexts` (list): An empty list, since this is a QA-only file.

        Raises:
            ValueError: If the file format is not supported (only .json and .jsonl are allowed) or if required fields are missing.

        Example:
            ```python
            documents = Dataset.load_dataset_qa("path/to/qa_dataset.jsonl")
            print(documents[0])
            # Output:
            # Document(question=Question("What is the capital of France?"), answers=Answer(["Paris"]), contexts=[])
            ```
        """
        documents = []

        # Load the file based on its extension (JSON or JSONL)
        if filepath.endswith(".json"):
            with open(filepath, encoding="utf-8") as file:
                data = json.load(file)
        elif filepath.endswith(".jsonl"):
            with open(filepath, encoding="utf-8") as file:
                data = [json.loads(line) for line in file]
        else:
            raise ValueError("Unsupported file format. Please use a JSON or JSONL file.")

        for entry in data:
            # Extract required fields
            #print(entry)
            question_text = entry.get("question")
            if not question_text:
                continue
                #raise ValueError("Each entry must have a 'question' field.")

            # Extract optional fields (answers and ID)

            answers = entry.get("answers", entry.get("answer", entry.get("golden_answers", [])))
            #print(answers)
            #asdadada
            if isinstance(answers, str):  
                answers = [answers]
            question_id = entry.get("id", None)

            # Create Question, Answer, and Document objects
            question_obj = Question(question_text)
            answer_obj = Answer(answers)
            document = Document(question=question_obj, answers=answer_obj, contexts=[])

            # Optionally store the ID in the document
            if question_id is not None:
                document.id = question_id

            documents.append(document)

        return documents

    def save_dataset(self,  output_path: str , save_reranked: bool= False, save_text:bool = False) -> None:
        """
        Saves the re-ranked documents in DPR format to a JSON or JSONL file.

        Args:
            output_path (str): The path to the output file (must be .json or .jsonl).

        Returns:
            None

        Example:
            ```python
            dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
            dataset.download()
            dataset.save_dataset("output.json")
            ```
        """
        dpr_data = []
        for doc in self.documents:
            if hasattr(doc.answers, "answers"):
                answers = doc.answers.answers
            elif isinstance(doc.answers, (list, tuple)):
                answers = list(doc.answers)
            elif isinstance(doc.answers, str):
                answers = [doc.answers]
            else:
                answers = [str(doc.answers)]

            dpr_entry = {
                "question": doc.question.question,
                "answers": answers,
                "ctxs": [ctx.to_dict(save_text) for ctx in doc.contexts]
            }
            if save_reranked:
                dpr_entry["reranked_ctxs"] = [ctx.to_dict(save_text) for ctx in doc.reorder_contexts]

            dpr_data.append(dpr_entry)

        if output_path.endswith(".json"):
            with open(output_path, 'w', encoding='utf-8') as file:
                json.dump(dpr_data, file, indent=4)
        elif output_path.endswith(".jsonl"):
            with open(output_path, 'w', encoding='utf-8') as file:
                for entry in dpr_data:
                    file.write(json.dumps(entry) + "\n")
        else:
            raise ValueError("Unsupported file format. Please use a .json or .jsonl file.")

    @staticmethod
    def save_documents(documents: List[Document], output_path: str, save_reranked: bool = False, save_text: bool = False) -> None:
        """
        Saves a list of Document objects in DPR format to a JSON or JSONL file.

        Args:
            documents (List[Document]): A list of Document objects containing questions, answers, and contexts.
            output_path (str): The path to save the DPR-formatted output (must be .json or .jsonl).
            save_reranked (bool, optional): Whether to save re-ranked contexts. Defaults to False.
            save_text (bool, optional): Whether to save the full text of the contexts. Defaults to False.

        Returns:
            None

        Example:
            ```python
            documents = [
                Document(
                    question=Question("What is the capital of France?"),
                    answers=Answer(["Paris"]),
                    contexts=[
                        Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris."),
                        Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany."),
                    ],
                )
            ]
            Dataset.save_documents(documents, "output.json", save_reranked=True, save_text=True)
            ```
        """
        dpr_data = []
        for doc in documents:
            dpr_entry = {
                "question": doc.question.question,
                "answers": doc.answers.answers,
                "ctxs": [ctx.to_dict(save_text) for ctx in doc.contexts]
            }
            if save_reranked and doc.reorder_contexts is not None:
                dpr_entry["reranked_ctxs"] = [ctx.to_dict(save_text) for ctx in doc.reorder_contexts]

            dpr_data.append(dpr_entry)

        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        if output_path.endswith(".json"):
            with open(output_path, "w", encoding="utf-8") as file:
                json.dump(dpr_data, file, indent=4)
        elif output_path.endswith(".jsonl"):
            with open(output_path, "w", encoding="utf-8") as file:
                for entry in dpr_data:
                    file.write(json.dumps(entry) + "\n")
        else:
            raise ValueError("Unsupported file format. Please use a .json or .jsonl file.")

        print(f"Saved {len(documents)} documents to {output_path}.")
    def __len__(self) -> int:
        """
        Returns the number of documents in the dataset.

        Returns:
            int: The number of documents in the dataset.
        """
        return len(self.documents)

    def __getitem__(self,idx) -> Document:
        """
        Retrieves the document at the specified index.

        Args:
            idx (int): The index of the document to retrieve.

        Returns:
            Document: The Document instance at the specified index.

        Raises:
            ValueError: If the dataset has not been loaded.

        Example:
            ```python
            dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
            dataset.download()
            document = dataset[0]
            print(document.question)  # Output: Question: What is the capital of France?
            ```
        """
        if not self.documents:
            raise ValueError("Dataset has not been loaded. Call `download()` to load the dataset.")
        return self.documents[idx]


    @staticmethod
    def available_dataset(file=None) -> None:
        """
        Prints information about available datasets.

        Example:
            ```python
            Dataset.available_dataset()
            ```
        """
        get_datasets_info(file=file)

__init__(retriever, dataset_name, n_docs=1000)

Initializes a Dataset instance.

Parameters:

Name Type Description Default
retriever str

The name of the retriever used to obtain the dataset.

 required
dataset_name str

The name of the dataset.

 required
n_docs int

The number of documents to include. Defaults to 1000.

 1000
Example 
dataset = Dataset(retriever='bm25', dataset_name='example_dataset', n_docs=500)
print(dataset.dataset_name)
Source code in rankify/dataset/dataset.py 
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
def __init__(self, retriever: str, dataset_name:str, n_docs:int = 1000) -> None:
    """
    Initializes a Dataset instance.

    Args:
        retriever (str): The name of the retriever used to obtain the dataset.
        dataset_name (str): The name of the dataset.
        n_docs (int, optional): The number of documents to include. Defaults to 1000.

    Example:
        ```python
        dataset = Dataset(retriever='bm25', dataset_name='example_dataset', n_docs=500)
        print(dataset.dataset_name)
        ```
    """
    self.dataset_name: str  = dataset_name
    self.retriever: str  = retriever
    self.n_docs: int = n_docs
    self.documents: Optional[List[Document]]  = None
    self.id_to_text_title: Optional[Dict] = None

update_contexts_from_passages()

Updates the text and title fields of contexts in all documents using the passages TSV file. If the TSV file is not already cached, it will be downloaded.

Raises:

Type Description
ValueError

If the dataset has not been loaded before calling this method.
Example 
dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
dataset.download()
dataset.update_contexts_from_passages()
Source code in rankify/dataset/dataset.py 
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
def update_contexts_from_passages(self) -> None:
    """
    Updates the text and title fields of contexts in all documents using the passages TSV file.
    If the TSV file is not already cached, it will be downloaded.

    Raises:
        ValueError: If the dataset has not been loaded before calling this method.

    Example:
        ```python
        dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
        dataset.download()
        dataset.update_contexts_from_passages()
        ```
    """
    if not self.documents:
        raise ValueError("Dataset has not been loaded. Call `download()` to load the dataset.")

    # Ensure TSV file is downloaded
    self._ensure_passages_downloaded()

    # Load the mapping of IDs to text and title
    id_to_text_title = self._load_passages_mapping()

    # Update contexts in all documents
    for document in tqdm(self.documents, desc="Retreive Text and Title for documents", unit="doc"):
        for context in document.contexts:
            if context.id in id_to_text_title:
                context_data = id_to_text_title[context.id]
                context.text = context_data["text"]
                context.title = context_data["title"]

download(force_download=True)

Downloads the dataset and loads it into memory.

Parameters:

Name Type Description Default
force_download bool

Whether to force downloading the dataset even if it already exists locally. Defaults to True.

 True

Returns:

Type Description
List[Document]

list[Document]: A list of Document instances loaded from the dataset.
Example 
dataset = Dataset(retriever='bm25', dataset_name='example_dataset', n_docs=500)
documents = dataset.download()
print(len(documents))
Source code in rankify/dataset/dataset.py 
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
def download(self, force_download:bool =True)-> List[Document]:
    """
    Downloads the dataset and loads it into memory.

    Args:
        force_download (bool, optional): Whether to force downloading the dataset even if it already exists locally. Defaults to True.

    Returns:
        list[Document]: A list of Document instances loaded from the dataset.

    Example:
        ```python
        dataset = Dataset(retriever='bm25', dataset_name='example_dataset', n_docs=500)
        documents = dataset.download()
        print(len(documents))
        ```
    """
    filepath= DownloadManger.download(self.retriever,self.dataset_name, force_download =force_download)
    self.documents= self.load_dataset(filepath, self.n_docs)
    if 'beir' not in self.dataset_name and 'dl' not in self.dataset_name:
        self.update_contexts_from_passages()
    return self.documents

load_dataset(filepath, n_docs=100) classmethod

Loads the dataset from a JSON file.

Parameters:

Name Type Description Default
filepath str

The path to the JSON file containing the dataset.

 required
n_docs int

The number of documents to load. Defaults to 100.

 100

Returns:

Type Description
List[Document]

list[Document]: A list of Document instances loaded from the JSON file.
Example 
filepath = 'example_dataset.json'
documents = Dataset.load_dataset(filepath, n_docs=50)
print(len(documents))
Source code in rankify/dataset/dataset.py 
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
@classmethod
def load_dataset(cls, filepath:str, n_docs: int= 100) -> List[Document]:
    """
    Loads the dataset from a JSON file.

    Args:
        filepath (str): The path to the JSON file containing the dataset.
        n_docs (int, optional): The number of documents to load. Defaults to 100.

    Returns:
        list[Document]: A list of Document instances loaded from the JSON file.

    Example:
        ```python
        filepath = 'example_dataset.json'
        documents = Dataset.load_dataset(filepath, n_docs=50)
        print(len(documents))
        ```
    """
    with open(filepath , encoding='utf-8') as file:
        data = json.load(file)
    data = [Document.from_dict(d,n_docs) for d in data]


    return data

load_dataset_qa(filepath) classmethod

Loads a QA dataset from a JSON or JSONL file. The dataset contains only questions, optionally answers, and optionally IDs.

Parameters:

Name Type Description Default
filepath str

The path to the JSON or JSONL file containing the dataset.

 required

Returns:

Type Description
List[Document]

List[Document]: A list of Document objects, each containing: - question (Question): The question object. - answers (Answer, optional): The answer object if available. - id (int, optional): The identifier if available. - contexts (list): An empty list, since this is a QA-only file.

Raises:

Type Description
ValueError

If the file format is not supported (only .json and .jsonl are allowed) or if required fields are missing.
Example 
documents = Dataset.load_dataset_qa("path/to/qa_dataset.jsonl")
print(documents[0])
# Output:
# Document(question=Question("What is the capital of France?"), answers=Answer(["Paris"]), contexts=[])
Source code in rankify/dataset/dataset.py 
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
@classmethod
def load_dataset_qa(cls, filepath: str) -> List[Document]:
    """
    Loads a QA dataset from a JSON or JSONL file. The dataset contains only questions, optionally answers, and optionally IDs.

    Args:
        filepath (str): The path to the JSON or JSONL file containing the dataset.

    Returns:
        List[Document]: A list of Document objects, each containing:
            - `question` (Question): The question object.
            - `answers` (Answer, optional): The answer object if available.
            - `id` (int, optional): The identifier if available.
            - `contexts` (list): An empty list, since this is a QA-only file.

    Raises:
        ValueError: If the file format is not supported (only .json and .jsonl are allowed) or if required fields are missing.

    Example:
        ```python
        documents = Dataset.load_dataset_qa("path/to/qa_dataset.jsonl")
        print(documents[0])
        # Output:
        # Document(question=Question("What is the capital of France?"), answers=Answer(["Paris"]), contexts=[])
        ```
    """
    documents = []

    # Load the file based on its extension (JSON or JSONL)
    if filepath.endswith(".json"):
        with open(filepath, encoding="utf-8") as file:
            data = json.load(file)
    elif filepath.endswith(".jsonl"):
        with open(filepath, encoding="utf-8") as file:
            data = [json.loads(line) for line in file]
    else:
        raise ValueError("Unsupported file format. Please use a JSON or JSONL file.")

    for entry in data:
        # Extract required fields
        #print(entry)
        question_text = entry.get("question")
        if not question_text:
            continue
            #raise ValueError("Each entry must have a 'question' field.")

        # Extract optional fields (answers and ID)

        answers = entry.get("answers", entry.get("answer", entry.get("golden_answers", [])))
        #print(answers)
        #asdadada
        if isinstance(answers, str):  
            answers = [answers]
        question_id = entry.get("id", None)

        # Create Question, Answer, and Document objects
        question_obj = Question(question_text)
        answer_obj = Answer(answers)
        document = Document(question=question_obj, answers=answer_obj, contexts=[])

        # Optionally store the ID in the document
        if question_id is not None:
            document.id = question_id

        documents.append(document)

    return documents

save_dataset(output_path, save_reranked=False, save_text=False)

Saves the re-ranked documents in DPR format to a JSON or JSONL file.

Parameters:

Name Type Description Default
output_path str

The path to the output file (must be .json or .jsonl).

 required

Returns:

Type Description
None

None
Example 
dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
dataset.download()
dataset.save_dataset("output.json")
Source code in rankify/dataset/dataset.py 
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
def save_dataset(self,  output_path: str , save_reranked: bool= False, save_text:bool = False) -> None:
    """
    Saves the re-ranked documents in DPR format to a JSON or JSONL file.

    Args:
        output_path (str): The path to the output file (must be .json or .jsonl).

    Returns:
        None

    Example:
        ```python
        dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
        dataset.download()
        dataset.save_dataset("output.json")
        ```
    """
    dpr_data = []
    for doc in self.documents:
        if hasattr(doc.answers, "answers"):
            answers = doc.answers.answers
        elif isinstance(doc.answers, (list, tuple)):
            answers = list(doc.answers)
        elif isinstance(doc.answers, str):
            answers = [doc.answers]
        else:
            answers = [str(doc.answers)]

        dpr_entry = {
            "question": doc.question.question,
            "answers": answers,
            "ctxs": [ctx.to_dict(save_text) for ctx in doc.contexts]
        }
        if save_reranked:
            dpr_entry["reranked_ctxs"] = [ctx.to_dict(save_text) for ctx in doc.reorder_contexts]

        dpr_data.append(dpr_entry)

    if output_path.endswith(".json"):
        with open(output_path, 'w', encoding='utf-8') as file:
            json.dump(dpr_data, file, indent=4)
    elif output_path.endswith(".jsonl"):
        with open(output_path, 'w', encoding='utf-8') as file:
            for entry in dpr_data:
                file.write(json.dumps(entry) + "\n")
    else:
        raise ValueError("Unsupported file format. Please use a .json or .jsonl file.")

save_documents(documents, output_path, save_reranked=False, save_text=False) staticmethod

Saves a list of Document objects in DPR format to a JSON or JSONL file.

Parameters:

Name Type Description Default
documents List[Document]

A list of Document objects containing questions, answers, and contexts.

 required
output_path str

The path to save the DPR-formatted output (must be .json or .jsonl).

 required
save_reranked bool

Whether to save re-ranked contexts. Defaults to False.

 False
save_text bool

Whether to save the full text of the contexts. Defaults to False.

 False

Returns:

Type Description
None

None
Example 
documents = [
    Document(
        question=Question("What is the capital of France?"),
        answers=Answer(["Paris"]),
        contexts=[
            Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris."),
            Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany."),
        ],
    )
]
Dataset.save_documents(documents, "output.json", save_reranked=True, save_text=True)
Source code in rankify/dataset/dataset.py 
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
@staticmethod
def save_documents(documents: List[Document], output_path: str, save_reranked: bool = False, save_text: bool = False) -> None:
    """
    Saves a list of Document objects in DPR format to a JSON or JSONL file.

    Args:
        documents (List[Document]): A list of Document objects containing questions, answers, and contexts.
        output_path (str): The path to save the DPR-formatted output (must be .json or .jsonl).
        save_reranked (bool, optional): Whether to save re-ranked contexts. Defaults to False.
        save_text (bool, optional): Whether to save the full text of the contexts. Defaults to False.

    Returns:
        None

    Example:
        ```python
        documents = [
            Document(
                question=Question("What is the capital of France?"),
                answers=Answer(["Paris"]),
                contexts=[
                    Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris."),
                    Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany."),
                ],
            )
        ]
        Dataset.save_documents(documents, "output.json", save_reranked=True, save_text=True)
        ```
    """
    dpr_data = []
    for doc in documents:
        dpr_entry = {
            "question": doc.question.question,
            "answers": doc.answers.answers,
            "ctxs": [ctx.to_dict(save_text) for ctx in doc.contexts]
        }
        if save_reranked and doc.reorder_contexts is not None:
            dpr_entry["reranked_ctxs"] = [ctx.to_dict(save_text) for ctx in doc.reorder_contexts]

        dpr_data.append(dpr_entry)

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    if output_path.endswith(".json"):
        with open(output_path, "w", encoding="utf-8") as file:
            json.dump(dpr_data, file, indent=4)
    elif output_path.endswith(".jsonl"):
        with open(output_path, "w", encoding="utf-8") as file:
            for entry in dpr_data:
                file.write(json.dumps(entry) + "\n")
    else:
        raise ValueError("Unsupported file format. Please use a .json or .jsonl file.")

    print(f"Saved {len(documents)} documents to {output_path}.")

__len__()

Returns the number of documents in the dataset.

Returns:

Name Type Description
int int

The number of documents in the dataset.
Source code in rankify/dataset/dataset.py 
624
625
626
627
628
629
630
631
def __len__(self) -> int:
    """
    Returns the number of documents in the dataset.

    Returns:
        int: The number of documents in the dataset.
    """
    return len(self.documents)

__getitem__(idx)

Retrieves the document at the specified index.

Parameters:

Name Type Description Default
idx int

The index of the document to retrieve.

 required

Returns:

Name Type Description
Document Document

The Document instance at the specified index.

Raises:

Type Description
ValueError

If the dataset has not been loaded.
Example 
dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
dataset.download()
document = dataset[0]
print(document.question)  # Output: Question: What is the capital of France?
Source code in rankify/dataset/dataset.py 
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
def __getitem__(self,idx) -> Document:
    """
    Retrieves the document at the specified index.

    Args:
        idx (int): The index of the document to retrieve.

    Returns:
        Document: The Document instance at the specified index.

    Raises:
        ValueError: If the dataset has not been loaded.

    Example:
        ```python
        dataset = Dataset(retriever="bm25", dataset_name="example_dataset", n_docs=500)
        dataset.download()
        document = dataset[0]
        print(document.question)  # Output: Question: What is the capital of France?
        ```
    """
    if not self.documents:
        raise ValueError("Dataset has not been loaded. Call `download()` to load the dataset.")
    return self.documents[idx]

available_dataset(file=None) staticmethod

Prints information about available datasets.

Example 
Dataset.available_dataset()
Source code in rankify/dataset/dataset.py 
659
660
661
662
663
664
665
666
667
668
669
@staticmethod
def available_dataset(file=None) -> None:
    """
    Prints information about available datasets.

    Example:
        ```python
        Dataset.available_dataset()
        ```
    """
    get_datasets_info(file=file)

get_datasets_info(file=None)

Source code in rankify/utils/dataset/utils.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
def get_datasets_info(file=None):
    table = PrettyTable(['Retriever', 'Dataset', 'Original ext', 'Compressed','Desc','URL'])
    for retriever, datasets in HF_PRE_DEFIND_DATASET.items():
        for dataset_name, dataset_info in datasets.items():

            flattened_entry = {
                'retriever': retriever,
                'dataset': dataset_name,
                'original_ext': dataset_info.get('original_ext'),
                'compressed': dataset_info.get('compressed'),
                'desc': dataset_info.get('desc'),
                'url': dataset_info.get('url')
            }
            table.add_row(flattened_entry.values())

    print(table, file=file)