trulens-benchmark 1.0.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- trulens_benchmark-1.0.1/PKG-INFO +24 -0
- trulens_benchmark-1.0.1/README.md +1 -0
- trulens_benchmark-1.0.1/pyproject.toml +34 -0
- trulens_benchmark-1.0.1/trulens/benchmark/__init__.py +15 -0
- trulens_benchmark-1.0.1/trulens/benchmark/answer_relevance_benchmark_small.ipynb +255 -0
- trulens_benchmark-1.0.1/trulens/benchmark/benchmark_frameworks/__init__.py +0 -0
- trulens_benchmark-1.0.1/trulens/benchmark/benchmark_frameworks/dataset/beir_loader.py +319 -0
- trulens_benchmark-1.0.1/trulens/benchmark/benchmark_frameworks/meta_evaluation_examples.ipynb +319 -0
- trulens_benchmark-1.0.1/trulens/benchmark/benchmark_frameworks/tru_benchmark_experiment.py +215 -0
- trulens_benchmark-1.0.1/trulens/benchmark/comprehensiveness_benchmark.ipynb +316 -0
- trulens_benchmark-1.0.1/trulens/benchmark/context_relevance_benchmark.ipynb +340 -0
- trulens_benchmark-1.0.1/trulens/benchmark/context_relevance_benchmark_calibration.ipynb +463 -0
- trulens_benchmark-1.0.1/trulens/benchmark/context_relevance_benchmark_small.ipynb +243 -0
- trulens_benchmark-1.0.1/trulens/benchmark/generate/__init__.py +0 -0
- trulens_benchmark-1.0.1/trulens/benchmark/generate/generate_test_set.py +208 -0
- trulens_benchmark-1.0.1/trulens/benchmark/groundedness_benchmark.ipynb +238 -0
- trulens_benchmark-1.0.1/trulens/benchmark/groundedness_benchmark_abstentions.ipynb +680 -0
- trulens_benchmark-1.0.1/trulens/benchmark/groundedness_benchmark_snowflake_arctic.ipynb +258 -0
- trulens_benchmark-1.0.1/trulens/benchmark/test_cases.py +291 -0
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: trulens-benchmark
|
|
3
|
+
Version: 1.0.1
|
|
4
|
+
Summary: Library to systematically track and evaluate LLM based applications.
|
|
5
|
+
Home-page: https://trulens.org/
|
|
6
|
+
License: MIT
|
|
7
|
+
Author: Snowflake Inc.
|
|
8
|
+
Author-email: ml-observability-wg-dl@snowflake.com
|
|
9
|
+
Requires-Python: >=3.9,<4.0
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Operating System :: OS Independent
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Requires-Dist: trulens-core (>=1.0.0,<2.0.0)
|
|
19
|
+
Project-URL: Documentation, https://trulens.org/trulens/getting_started/
|
|
20
|
+
Project-URL: Repository, https://github.com/truera/trulens
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# trulens-benchmark
|
|
24
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# trulens-benchmark
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
build-backend = "poetry.core.masonry.api"
|
|
3
|
+
requires = [
|
|
4
|
+
"poetry-core",
|
|
5
|
+
]
|
|
6
|
+
|
|
7
|
+
[tool.poetry]
|
|
8
|
+
name = "trulens-benchmark"
|
|
9
|
+
version = "1.0.1"
|
|
10
|
+
description = "Library to systematically track and evaluate LLM based applications."
|
|
11
|
+
authors = [
|
|
12
|
+
"Snowflake Inc. <ml-observability-wg-dl@snowflake.com>",
|
|
13
|
+
]
|
|
14
|
+
license = "MIT"
|
|
15
|
+
readme = "README.md"
|
|
16
|
+
packages = [
|
|
17
|
+
{ include = "trulens" },
|
|
18
|
+
]
|
|
19
|
+
homepage = "https://trulens.org/"
|
|
20
|
+
documentation = "https://trulens.org/trulens/getting_started/"
|
|
21
|
+
repository = "https://github.com/truera/trulens"
|
|
22
|
+
classifiers = [
|
|
23
|
+
"Programming Language :: Python :: 3",
|
|
24
|
+
"Operating System :: OS Independent",
|
|
25
|
+
"Development Status :: 3 - Alpha",
|
|
26
|
+
"License :: OSI Approved :: MIT License",
|
|
27
|
+
]
|
|
28
|
+
|
|
29
|
+
[tool.poetry.dependencies]
|
|
30
|
+
python = "^3.9"
|
|
31
|
+
trulens-core = { version = "^1.0.0" }
|
|
32
|
+
|
|
33
|
+
[tool.poetry.group.dev.dependencies]
|
|
34
|
+
trulens-core = { path = "../core", develop = true }
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
"""
|
|
2
|
+
!!! note "Additional Dependency Required"
|
|
3
|
+
|
|
4
|
+
To use this module, you must have the `trulens-benchmark` package installed.
|
|
5
|
+
|
|
6
|
+
```bash
|
|
7
|
+
pip install trulens-benchmark
|
|
8
|
+
```
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from importlib.metadata import version
|
|
12
|
+
|
|
13
|
+
from trulens.core.utils.imports import safe_importlib_package_name
|
|
14
|
+
|
|
15
|
+
__version__ = version(safe_importlib_package_name(__package__ or __name__))
|
|
@@ -0,0 +1,255 @@
|
|
|
1
|
+
{
|
|
2
|
+
"cells": [
|
|
3
|
+
{
|
|
4
|
+
"attachments": {},
|
|
5
|
+
"cell_type": "markdown",
|
|
6
|
+
"metadata": {},
|
|
7
|
+
"source": [
|
|
8
|
+
"# 📓 Answer Relevance Feedback Evaluation\n",
|
|
9
|
+
"In many ways, feedbacks can be thought of as LLM apps themselves. Given text,\n",
|
|
10
|
+
"they return some result. Thinking in this way, we can use _TruLens_ to evaluate\n",
|
|
11
|
+
"and track our feedback quality. We can even do this for different models (e.g.\n",
|
|
12
|
+
"gpt-3.5 and gpt-4) or prompting schemes (such as chain-of-thought reasoning).\n",
|
|
13
|
+
"\n",
|
|
14
|
+
"This notebook follows an evaluation of a set of test cases. You are encouraged\n",
|
|
15
|
+
"to run this on your own and even expand the test cases to evaluate performance\n",
|
|
16
|
+
"on test cases applicable to your scenario or domain."
|
|
17
|
+
]
|
|
18
|
+
},
|
|
19
|
+
{
|
|
20
|
+
"cell_type": "code",
|
|
21
|
+
"execution_count": null,
|
|
22
|
+
"metadata": {},
|
|
23
|
+
"outputs": [],
|
|
24
|
+
"source": [
|
|
25
|
+
"# Import relevance feedback function\n",
|
|
26
|
+
"from test_cases import answer_relevance_golden_set\n",
|
|
27
|
+
"from trulens.apps.basic import TruBasicApp\n",
|
|
28
|
+
"from trulens.core import Feedback\n",
|
|
29
|
+
"from trulens.core import Select\n",
|
|
30
|
+
"from trulens.core import TruSession\n",
|
|
31
|
+
"from trulens.feedback import GroundTruthAgreement\n",
|
|
32
|
+
"from trulens.providers.litellm import LiteLLM\n",
|
|
33
|
+
"from trulens.providers.openai import OpenAI\n",
|
|
34
|
+
"\n",
|
|
35
|
+
"TruSession().reset_database()"
|
|
36
|
+
]
|
|
37
|
+
},
|
|
38
|
+
{
|
|
39
|
+
"cell_type": "code",
|
|
40
|
+
"execution_count": null,
|
|
41
|
+
"metadata": {},
|
|
42
|
+
"outputs": [],
|
|
43
|
+
"source": [
|
|
44
|
+
"import os\n",
|
|
45
|
+
"\n",
|
|
46
|
+
"os.environ[\"OPENAI_API_KEY\"] = \"...\"\n",
|
|
47
|
+
"os.environ[\"COHERE_API_KEY\"] = \"...\"\n",
|
|
48
|
+
"os.environ[\"HUGGINGFACE_API_KEY\"] = \"...\"\n",
|
|
49
|
+
"os.environ[\"ANTHROPIC_API_KEY\"] = \"...\"\n",
|
|
50
|
+
"os.environ[\"TOGETHERAI_API_KEY\"] = \"...\""
|
|
51
|
+
]
|
|
52
|
+
},
|
|
53
|
+
{
|
|
54
|
+
"cell_type": "code",
|
|
55
|
+
"execution_count": null,
|
|
56
|
+
"metadata": {},
|
|
57
|
+
"outputs": [],
|
|
58
|
+
"source": [
|
|
59
|
+
"# GPT 3.5\n",
|
|
60
|
+
"turbo = OpenAI(model_engine=\"gpt-3.5-turbo\")\n",
|
|
61
|
+
"\n",
|
|
62
|
+
"\n",
|
|
63
|
+
"def wrapped_relevance_turbo(input, output):\n",
|
|
64
|
+
" return turbo.relevance(input, output)\n",
|
|
65
|
+
"\n",
|
|
66
|
+
"\n",
|
|
67
|
+
"# GPT 4\n",
|
|
68
|
+
"gpt4 = OpenAI(model_engine=\"gpt-4\")\n",
|
|
69
|
+
"\n",
|
|
70
|
+
"\n",
|
|
71
|
+
"def wrapped_relevance_gpt4(input, output):\n",
|
|
72
|
+
" return gpt4.relevance(input, output)\n",
|
|
73
|
+
"\n",
|
|
74
|
+
"\n",
|
|
75
|
+
"# Cohere\n",
|
|
76
|
+
"command_nightly = LiteLLM(model_engine=\"cohere/command-nightly\")\n",
|
|
77
|
+
"\n",
|
|
78
|
+
"\n",
|
|
79
|
+
"def wrapped_relevance_command_nightly(input, output):\n",
|
|
80
|
+
" return command_nightly.relevance(input, output)\n",
|
|
81
|
+
"\n",
|
|
82
|
+
"\n",
|
|
83
|
+
"# Anthropic\n",
|
|
84
|
+
"claude_1 = LiteLLM(model_engine=\"claude-instant-1\")\n",
|
|
85
|
+
"\n",
|
|
86
|
+
"\n",
|
|
87
|
+
"def wrapped_relevance_claude1(input, output):\n",
|
|
88
|
+
" return claude_1.relevance(input, output)\n",
|
|
89
|
+
"\n",
|
|
90
|
+
"\n",
|
|
91
|
+
"claude_2 = LiteLLM(model_engine=\"claude-2\")\n",
|
|
92
|
+
"\n",
|
|
93
|
+
"\n",
|
|
94
|
+
"def wrapped_relevance_claude2(input, output):\n",
|
|
95
|
+
" return claude_2.relevance(input, output)\n",
|
|
96
|
+
"\n",
|
|
97
|
+
"\n",
|
|
98
|
+
"# Meta\n",
|
|
99
|
+
"llama_2_13b = LiteLLM(\n",
|
|
100
|
+
" model_engine=\"together_ai/togethercomputer/Llama-2-7B-32K-Instruct\"\n",
|
|
101
|
+
")\n",
|
|
102
|
+
"\n",
|
|
103
|
+
"\n",
|
|
104
|
+
"def wrapped_relevance_llama2(input, output):\n",
|
|
105
|
+
" return llama_2_13b.relevance(input, output)"
|
|
106
|
+
]
|
|
107
|
+
},
|
|
108
|
+
{
|
|
109
|
+
"attachments": {},
|
|
110
|
+
"cell_type": "markdown",
|
|
111
|
+
"metadata": {},
|
|
112
|
+
"source": [
|
|
113
|
+
"Here we'll set up our golden set as a set of prompts, responses and expected\n",
|
|
114
|
+
"scores stored in `test_cases.py`. Then, our numeric_difference method will look\n",
|
|
115
|
+
"up the expected score for each prompt/response pair by **exact match**. After\n",
|
|
116
|
+
"looking up the expected score, we will then take the L1 difference between the\n",
|
|
117
|
+
"actual score and expected score."
|
|
118
|
+
]
|
|
119
|
+
},
|
|
120
|
+
{
|
|
121
|
+
"cell_type": "code",
|
|
122
|
+
"execution_count": null,
|
|
123
|
+
"metadata": {},
|
|
124
|
+
"outputs": [],
|
|
125
|
+
"source": [
|
|
126
|
+
"# Create a Feedback object using the numeric_difference method of the\n",
|
|
127
|
+
"# ground_truth object\n",
|
|
128
|
+
"ground_truth = GroundTruthAgreement(\n",
|
|
129
|
+
" answer_relevance_golden_set, provider=OpenAI()\n",
|
|
130
|
+
")\n",
|
|
131
|
+
"\n",
|
|
132
|
+
"# Call the numeric_difference method with app and record and aggregate to get\n",
|
|
133
|
+
"# the mean absolute error\n",
|
|
134
|
+
"f_mae = (\n",
|
|
135
|
+
" Feedback(ground_truth.mae, name=\"Mean Absolute Error\")\n",
|
|
136
|
+
" .on(Select.Record.calls[0].args.args[0])\n",
|
|
137
|
+
" .on(Select.Record.calls[0].args.args[1])\n",
|
|
138
|
+
" .on_output()\n",
|
|
139
|
+
")"
|
|
140
|
+
]
|
|
141
|
+
},
|
|
142
|
+
{
|
|
143
|
+
"cell_type": "code",
|
|
144
|
+
"execution_count": null,
|
|
145
|
+
"metadata": {},
|
|
146
|
+
"outputs": [],
|
|
147
|
+
"source": [
|
|
148
|
+
"tru_wrapped_relevance_turbo = TruBasicApp(\n",
|
|
149
|
+
" wrapped_relevance_turbo,\n",
|
|
150
|
+
" app_name=\"answer relevance\",\n",
|
|
151
|
+
" app_version=\"gpt-3.5-turbo\",\n",
|
|
152
|
+
" feedbacks=[f_mae],\n",
|
|
153
|
+
")\n",
|
|
154
|
+
"\n",
|
|
155
|
+
"tru_wrapped_relevance_gpt4 = TruBasicApp(\n",
|
|
156
|
+
" wrapped_relevance_gpt4,\n",
|
|
157
|
+
" app_name=\"answer relevance\",\n",
|
|
158
|
+
" app_version=\"gpt-4\",\n",
|
|
159
|
+
" feedbacks=[f_mae],\n",
|
|
160
|
+
")\n",
|
|
161
|
+
"\n",
|
|
162
|
+
"tru_wrapped_relevance_commandnightly = TruBasicApp(\n",
|
|
163
|
+
" wrapped_relevance_command_nightly,\n",
|
|
164
|
+
" app_name=\"answer relevance\",\n",
|
|
165
|
+
" app_version=\"Command-Nightly\",\n",
|
|
166
|
+
" feedbacks=[f_mae],\n",
|
|
167
|
+
")\n",
|
|
168
|
+
"\n",
|
|
169
|
+
"tru_wrapped_relevance_claude1 = TruBasicApp(\n",
|
|
170
|
+
" wrapped_relevance_claude1,\n",
|
|
171
|
+
" app_name=\"answer relevance\",\n",
|
|
172
|
+
" app_version=\"Claude 1\",\n",
|
|
173
|
+
" feedbacks=[f_mae],\n",
|
|
174
|
+
")\n",
|
|
175
|
+
"\n",
|
|
176
|
+
"tru_wrapped_relevance_claude2 = TruBasicApp(\n",
|
|
177
|
+
" wrapped_relevance_claude2,\n",
|
|
178
|
+
" app_name=\"answer relevance\",\n",
|
|
179
|
+
" app_version=\"Claude 2\",\n",
|
|
180
|
+
" feedbacks=[f_mae],\n",
|
|
181
|
+
")\n",
|
|
182
|
+
"\n",
|
|
183
|
+
"tru_wrapped_relevance_llama2 = TruBasicApp(\n",
|
|
184
|
+
" wrapped_relevance_llama2,\n",
|
|
185
|
+
" app_name=\"answer relevance\",\n",
|
|
186
|
+
" app_version=\"Llama-2-13b\",\n",
|
|
187
|
+
" feedbacks=[f_mae],\n",
|
|
188
|
+
")"
|
|
189
|
+
]
|
|
190
|
+
},
|
|
191
|
+
{
|
|
192
|
+
"cell_type": "code",
|
|
193
|
+
"execution_count": null,
|
|
194
|
+
"metadata": {},
|
|
195
|
+
"outputs": [],
|
|
196
|
+
"source": [
|
|
197
|
+
"for i in range(len(answer_relevance_golden_set)):\n",
|
|
198
|
+
" prompt = answer_relevance_golden_set[i][\"query\"]\n",
|
|
199
|
+
" response = answer_relevance_golden_set[i][\"response\"]\n",
|
|
200
|
+
"\n",
|
|
201
|
+
" with tru_wrapped_relevance_turbo as recording:\n",
|
|
202
|
+
" tru_wrapped_relevance_turbo.app(prompt, response)\n",
|
|
203
|
+
"\n",
|
|
204
|
+
" with tru_wrapped_relevance_gpt4 as recording:\n",
|
|
205
|
+
" tru_wrapped_relevance_gpt4.app(prompt, response)\n",
|
|
206
|
+
"\n",
|
|
207
|
+
" with tru_wrapped_relevance_commandnightly as recording:\n",
|
|
208
|
+
" tru_wrapped_relevance_commandnightly.app(prompt, response)\n",
|
|
209
|
+
"\n",
|
|
210
|
+
" with tru_wrapped_relevance_claude1 as recording:\n",
|
|
211
|
+
" tru_wrapped_relevance_claude1.app(prompt, response)\n",
|
|
212
|
+
"\n",
|
|
213
|
+
" with tru_wrapped_relevance_claude2 as recording:\n",
|
|
214
|
+
" tru_wrapped_relevance_claude2.app(prompt, response)\n",
|
|
215
|
+
"\n",
|
|
216
|
+
" with tru_wrapped_relevance_llama2 as recording:\n",
|
|
217
|
+
" tru_wrapped_relevance_llama2.app(prompt, response)"
|
|
218
|
+
]
|
|
219
|
+
},
|
|
220
|
+
{
|
|
221
|
+
"cell_type": "code",
|
|
222
|
+
"execution_count": null,
|
|
223
|
+
"metadata": {},
|
|
224
|
+
"outputs": [],
|
|
225
|
+
"source": [
|
|
226
|
+
"TruSession().get_leaderboard().sort_values(by=\"Mean Absolute Error\")"
|
|
227
|
+
]
|
|
228
|
+
}
|
|
229
|
+
],
|
|
230
|
+
"metadata": {
|
|
231
|
+
"kernelspec": {
|
|
232
|
+
"display_name": "Python 3.11.4 ('agents')",
|
|
233
|
+
"language": "python",
|
|
234
|
+
"name": "python3"
|
|
235
|
+
},
|
|
236
|
+
"language_info": {
|
|
237
|
+
"codemirror_mode": {
|
|
238
|
+
"name": "ipython",
|
|
239
|
+
"version": 3
|
|
240
|
+
},
|
|
241
|
+
"file_extension": ".py",
|
|
242
|
+
"mimetype": "text/x-python",
|
|
243
|
+
"name": "python",
|
|
244
|
+
"nbconvert_exporter": "python",
|
|
245
|
+
"pygments_lexer": "ipython3"
|
|
246
|
+
},
|
|
247
|
+
"vscode": {
|
|
248
|
+
"interpreter": {
|
|
249
|
+
"hash": "7d153714b979d5e6d08dd8ec90712dd93bff2c9b6c1f0c118169738af3430cd4"
|
|
250
|
+
}
|
|
251
|
+
}
|
|
252
|
+
},
|
|
253
|
+
"nbformat": 4,
|
|
254
|
+
"nbformat_minor": 2
|
|
255
|
+
}
|
|
File without changes
|
|
@@ -0,0 +1,319 @@
|
|
|
1
|
+
import csv
|
|
2
|
+
import json
|
|
3
|
+
import logging
|
|
4
|
+
import os
|
|
5
|
+
from typing import Any, Dict, Generator, List, Optional, Tuple
|
|
6
|
+
import zipfile
|
|
7
|
+
|
|
8
|
+
import pandas as pd
|
|
9
|
+
import requests
|
|
10
|
+
from tqdm.autonotebook import tqdm
|
|
11
|
+
from trulens.core import TruSession
|
|
12
|
+
|
|
13
|
+
logger = logging.getLogger(__name__)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def download_and_unzip(url: str, out_dir: str, chunk_size: int = 1024) -> str:
|
|
17
|
+
def download_url(url: str, save_path: str, chunk_size: int = 1024):
|
|
18
|
+
"""Download url with progress bar using tqdm
|
|
19
|
+
https://stackoverflow.com/questions/15644964/python-progress-bar-and-downloads
|
|
20
|
+
|
|
21
|
+
Args:
|
|
22
|
+
url (str): downloadable url
|
|
23
|
+
save_path (str): local path to save the downloaded file
|
|
24
|
+
chunk_size (int, optional): chunking of files. Defaults to 1024.
|
|
25
|
+
"""
|
|
26
|
+
r = requests.get(url, stream=True)
|
|
27
|
+
total = int(r.headers.get("Content-Length", 0))
|
|
28
|
+
with open(save_path, "wb") as fd, tqdm(
|
|
29
|
+
desc=save_path,
|
|
30
|
+
total=total,
|
|
31
|
+
unit="iB",
|
|
32
|
+
unit_scale=True,
|
|
33
|
+
unit_divisor=chunk_size,
|
|
34
|
+
) as bar:
|
|
35
|
+
for data in r.iter_content(chunk_size=chunk_size):
|
|
36
|
+
size = fd.write(data)
|
|
37
|
+
bar.update(size)
|
|
38
|
+
|
|
39
|
+
def unzip(zip_file: str, out_dir: str):
|
|
40
|
+
zip_ = zipfile.ZipFile(zip_file, "r")
|
|
41
|
+
zip_.extractall(path=out_dir)
|
|
42
|
+
zip_.close()
|
|
43
|
+
|
|
44
|
+
os.makedirs(out_dir, exist_ok=True)
|
|
45
|
+
dataset = url.split("/")[-1]
|
|
46
|
+
zip_file = os.path.join(out_dir, dataset)
|
|
47
|
+
|
|
48
|
+
if not os.path.isfile(zip_file):
|
|
49
|
+
logger.info("Downloading {} ...".format(dataset))
|
|
50
|
+
download_url(url, zip_file, chunk_size)
|
|
51
|
+
|
|
52
|
+
if not os.path.isdir(zip_file.replace(".zip", "")):
|
|
53
|
+
logger.info("Unzipping {} ...".format(dataset))
|
|
54
|
+
unzip(zip_file, out_dir)
|
|
55
|
+
|
|
56
|
+
return os.path.join(out_dir, dataset.replace(".zip", ""))
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
class TruBEIRDataLoader:
|
|
60
|
+
def __init__(
|
|
61
|
+
self,
|
|
62
|
+
data_folder: str,
|
|
63
|
+
dataset_name: str,
|
|
64
|
+
corpus_file: str = "corpus.jsonl",
|
|
65
|
+
query_file: str = "queries.jsonl",
|
|
66
|
+
qrels_folder: str = "qrels",
|
|
67
|
+
qrels_file: str = "",
|
|
68
|
+
):
|
|
69
|
+
"""
|
|
70
|
+
A utility class to load BEIR datasets into TruLens. Similar to https://github.com/beir-cellar/beir but with slightly
|
|
71
|
+
more efficient implementation of processing and loading of the datasets using generators.
|
|
72
|
+
|
|
73
|
+
Args:
|
|
74
|
+
data_folder (str): local / remote path to store the downloaded dataset.
|
|
75
|
+
dataset_name (str): Name of the dataset to be loaded. Must be one of https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/
|
|
76
|
+
corpus_file (str, optional): file name of the corpus. Defaults to "corpus.jsonl".
|
|
77
|
+
query_file (str, optional): file name of all the queries. Defaults to "queries.jsonl".
|
|
78
|
+
qrels_folder (str, optional): folder name of qrels (relevance annotation). Defaults to "qrels".
|
|
79
|
+
qrels_file (str, optional): file name of qrels (relevance annotation). Defaults to "".
|
|
80
|
+
"""
|
|
81
|
+
self.qrels = {}
|
|
82
|
+
|
|
83
|
+
self.dataset_name = dataset_name
|
|
84
|
+
self.data_folder = data_folder
|
|
85
|
+
|
|
86
|
+
self.corpus_file = (
|
|
87
|
+
os.path.join(self.data_folder, dataset_name, corpus_file)
|
|
88
|
+
if self.data_folder
|
|
89
|
+
else corpus_file
|
|
90
|
+
)
|
|
91
|
+
self.query_file = (
|
|
92
|
+
os.path.join(self.data_folder, dataset_name, query_file)
|
|
93
|
+
if self.data_folder
|
|
94
|
+
else query_file
|
|
95
|
+
)
|
|
96
|
+
self.qrels_folder = os.path.join(
|
|
97
|
+
self.data_folder, dataset_name, qrels_folder
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
self.qrels_file = qrels_file
|
|
101
|
+
|
|
102
|
+
@staticmethod
|
|
103
|
+
def check(fIn: str, ext: str):
|
|
104
|
+
if not os.path.exists(fIn):
|
|
105
|
+
raise ValueError(
|
|
106
|
+
f"File {fIn} not present! Please provide accurate file."
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
if not fIn.endswith(ext):
|
|
110
|
+
raise ValueError(f"File {fIn} must be present with extension {ext}")
|
|
111
|
+
|
|
112
|
+
def _load_qrels(self):
|
|
113
|
+
reader = csv.reader(
|
|
114
|
+
open(self.qrels_file, encoding="utf-8"),
|
|
115
|
+
delimiter="\t",
|
|
116
|
+
quoting=csv.QUOTE_MINIMAL,
|
|
117
|
+
)
|
|
118
|
+
next(reader)
|
|
119
|
+
|
|
120
|
+
for id, row in enumerate(reader):
|
|
121
|
+
query_id, corpus_id, score = row[0], row[1], int(row[2])
|
|
122
|
+
|
|
123
|
+
if query_id not in self.qrels:
|
|
124
|
+
self.qrels[query_id] = {corpus_id: score}
|
|
125
|
+
else:
|
|
126
|
+
self.qrels[query_id][corpus_id] = score
|
|
127
|
+
|
|
128
|
+
def _load_generators(
|
|
129
|
+
self, split="test"
|
|
130
|
+
) -> Tuple[
|
|
131
|
+
Generator[Dict[str, Dict[str, str]], None, None],
|
|
132
|
+
Generator[Dict[str, str], None, None],
|
|
133
|
+
]:
|
|
134
|
+
"""
|
|
135
|
+
Load corpus, queries, and qrels as generators.
|
|
136
|
+
|
|
137
|
+
Args:
|
|
138
|
+
split (str, optional): Dataset split to load. Defaults to "test".
|
|
139
|
+
|
|
140
|
+
Returns:
|
|
141
|
+
Tuple of generators for corpus, queries, and qrels.
|
|
142
|
+
"""
|
|
143
|
+
self.qrels_file = os.path.join(self.qrels_folder, split + ".tsv")
|
|
144
|
+
self.check(fIn=self.corpus_file, ext="jsonl")
|
|
145
|
+
self.check(fIn=self.query_file, ext="jsonl")
|
|
146
|
+
self.check(fIn=self.qrels_file, ext="tsv")
|
|
147
|
+
|
|
148
|
+
corpus_gen = self._corpus_generator()
|
|
149
|
+
|
|
150
|
+
self._load_qrels()
|
|
151
|
+
queries_gen = self._queries_generator()
|
|
152
|
+
|
|
153
|
+
return corpus_gen, queries_gen
|
|
154
|
+
|
|
155
|
+
def _corpus_generator(
|
|
156
|
+
self,
|
|
157
|
+
) -> Generator[Dict[str, Dict[str, str]], None, None]:
|
|
158
|
+
"""
|
|
159
|
+
Generator to load corpus data incrementally.
|
|
160
|
+
"""
|
|
161
|
+
self.check(fIn=self.corpus_file, ext="jsonl")
|
|
162
|
+
|
|
163
|
+
with open(self.corpus_file, encoding="utf8") as fIn:
|
|
164
|
+
for line in fIn:
|
|
165
|
+
line = json.loads(line)
|
|
166
|
+
yield {
|
|
167
|
+
line.get("_id"): {
|
|
168
|
+
"text": line.get("text"),
|
|
169
|
+
"title": line.get("title"),
|
|
170
|
+
}
|
|
171
|
+
}
|
|
172
|
+
|
|
173
|
+
def _queries_generator(self) -> Generator[Dict[str, str], None, None]:
|
|
174
|
+
"""
|
|
175
|
+
Generator to load queries incrementally.
|
|
176
|
+
"""
|
|
177
|
+
self.check(fIn=self.query_file, ext="jsonl")
|
|
178
|
+
|
|
179
|
+
with open(self.query_file, encoding="utf8") as fIn:
|
|
180
|
+
for line in fIn:
|
|
181
|
+
line = json.loads(line)
|
|
182
|
+
if line.get("_id") in self.qrels:
|
|
183
|
+
yield {line.get("_id"): line.get("text")}
|
|
184
|
+
|
|
185
|
+
def _process_dataset(
|
|
186
|
+
self, split="test", chunk_size=None
|
|
187
|
+
) -> Generator[List[Dict[str, Any]], None, None]:
|
|
188
|
+
"""
|
|
189
|
+
Common method to process the BEIR dataset into entries.
|
|
190
|
+
This method handles downloading, loading, and processing the dataset.
|
|
191
|
+
|
|
192
|
+
Args:
|
|
193
|
+
split (str, optional): Dataset split to load. Defaults to "test".
|
|
194
|
+
chunk_size (int, optional): Number of records to process in each chunk. Defaults to None.
|
|
195
|
+
|
|
196
|
+
Yields:
|
|
197
|
+
List[Dict[str, Any]]: List of dataset entries (a chunk if chunk_size is specified).
|
|
198
|
+
"""
|
|
199
|
+
|
|
200
|
+
url = f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{self.dataset_name}.zip"
|
|
201
|
+
|
|
202
|
+
logger.info(f"Downloading {self.dataset_name} dataset from {url}")
|
|
203
|
+
download_and_unzip(url, self.data_folder)
|
|
204
|
+
|
|
205
|
+
corpus_gen, queries_gen = self._load_generators(split=split)
|
|
206
|
+
|
|
207
|
+
dataset_entries = []
|
|
208
|
+
|
|
209
|
+
corpus_dict = {}
|
|
210
|
+
for corpus in list(
|
|
211
|
+
corpus_gen
|
|
212
|
+
): # TODO: (Daniel) this is still going to be a memory hog for large datasets, but if we don't
|
|
213
|
+
# load the entire corpus into memory, we'll have to re-read the entire corpus using generator for each query
|
|
214
|
+
corpus_dict.update(corpus)
|
|
215
|
+
|
|
216
|
+
# Iterate over queries generator and process entries
|
|
217
|
+
for i, query in enumerate(queries_gen, 1):
|
|
218
|
+
for query_id, query_text in query.items():
|
|
219
|
+
doc_to_rel = self.qrels.get(query_id, {})
|
|
220
|
+
|
|
221
|
+
expected_chunks = []
|
|
222
|
+
|
|
223
|
+
for corpus_id in doc_to_rel.keys():
|
|
224
|
+
if corpus_id in corpus_dict:
|
|
225
|
+
corpus_entry = corpus_dict[corpus_id]
|
|
226
|
+
expected_chunks.append({
|
|
227
|
+
"text": corpus_entry["text"],
|
|
228
|
+
"title": corpus_entry.get("title"),
|
|
229
|
+
"expected_score": doc_to_rel.get(
|
|
230
|
+
corpus_id, 0
|
|
231
|
+
), # Default relevance score to 0 if not found
|
|
232
|
+
})
|
|
233
|
+
|
|
234
|
+
dataset_entries.append({
|
|
235
|
+
"query_id": query_id,
|
|
236
|
+
"query": query_text,
|
|
237
|
+
"expected_response": None, # expected response can be empty for IR datasets
|
|
238
|
+
"expected_chunks": expected_chunks,
|
|
239
|
+
"meta": {"source": self.dataset_name},
|
|
240
|
+
})
|
|
241
|
+
|
|
242
|
+
# Yield the chunk if chunk_size is specified
|
|
243
|
+
if chunk_size and i % chunk_size == 0:
|
|
244
|
+
yield dataset_entries
|
|
245
|
+
dataset_entries = []
|
|
246
|
+
|
|
247
|
+
if dataset_entries:
|
|
248
|
+
yield dataset_entries
|
|
249
|
+
|
|
250
|
+
def load_dataset_to_df(
|
|
251
|
+
self,
|
|
252
|
+
split="test",
|
|
253
|
+
download=True,
|
|
254
|
+
) -> pd.DataFrame:
|
|
255
|
+
"""
|
|
256
|
+
load BEIR dataset into dataframe with pre-processed fields to match expected TruLens schemas.
|
|
257
|
+
Note this method loads the entire dataset into memory at once.
|
|
258
|
+
Args:
|
|
259
|
+
split (str, optional): Defaults to "test".
|
|
260
|
+
download (bool, optional): If False, remove the downloaded dataset file after processing. Defaults to True.
|
|
261
|
+
Returns:
|
|
262
|
+
pd.DataFrame: DataFrame with the BEIR dataset
|
|
263
|
+
"""
|
|
264
|
+
dataset_entries = []
|
|
265
|
+
for chunk in self._process_dataset(split=split):
|
|
266
|
+
dataset_entries.extend(chunk)
|
|
267
|
+
|
|
268
|
+
if not download:
|
|
269
|
+
logger.info(f"Cleaning up downloaded {self.dataset_name} dataset")
|
|
270
|
+
os.system(
|
|
271
|
+
f"rm -rf {os.path.join(self.data_folder, self.dataset_name)}"
|
|
272
|
+
)
|
|
273
|
+
os.system(
|
|
274
|
+
f"rm -rf {os.path.join(self.data_folder, self.dataset_name)}.zip"
|
|
275
|
+
)
|
|
276
|
+
|
|
277
|
+
return pd.DataFrame(dataset_entries)
|
|
278
|
+
|
|
279
|
+
def persist_dataset(
|
|
280
|
+
self,
|
|
281
|
+
session: TruSession,
|
|
282
|
+
dataset_name: str,
|
|
283
|
+
dataset_metadata: Optional[Dict[str, Any]] = None,
|
|
284
|
+
split="test",
|
|
285
|
+
download=True,
|
|
286
|
+
chunk_size=1000,
|
|
287
|
+
):
|
|
288
|
+
"""
|
|
289
|
+
Persist BEIR dataset into DB with pre-processed fields to match expected TruLens schemas.
|
|
290
|
+
Note this method handle chunking of the dataset to avoid loading the entire dataset into memory at once by default.
|
|
291
|
+
Args:
|
|
292
|
+
split (str, optional): Defaults to "test".
|
|
293
|
+
session (TruSession): TruSession instance to persist the dataset.
|
|
294
|
+
dataset_name (str): Name of the dataset to be persisted - Note this can be different from the standardized BEIR dataset names.
|
|
295
|
+
dataset_metadata (Optional[Dict[str, Any]], optional): Metadata for the dataset.
|
|
296
|
+
download (bool, optional): If False, remove the downloaded dataset file after processing. Defaults to True.
|
|
297
|
+
Returns:
|
|
298
|
+
pd.DataFrame: DataFrame with the BEIR dataset
|
|
299
|
+
"""
|
|
300
|
+
for chunk in self._process_dataset(split=split, chunk_size=chunk_size):
|
|
301
|
+
df_chunk = pd.DataFrame(chunk)
|
|
302
|
+
session.add_ground_truth_to_dataset(
|
|
303
|
+
dataset_name=dataset_name,
|
|
304
|
+
ground_truth_df=df_chunk,
|
|
305
|
+
dataset_metadata=dataset_metadata,
|
|
306
|
+
)
|
|
307
|
+
|
|
308
|
+
if not download:
|
|
309
|
+
logger.info(f"Cleaning up downloaded {self.dataset_name} dataset")
|
|
310
|
+
os.system(
|
|
311
|
+
f"rm -rf {os.path.join(self.data_folder, self.dataset_name)}"
|
|
312
|
+
)
|
|
313
|
+
os.system(
|
|
314
|
+
f"rm -rf {os.path.join(self.data_folder, self.dataset_name)}.zip"
|
|
315
|
+
)
|
|
316
|
+
|
|
317
|
+
logger.info(
|
|
318
|
+
f"Finished processing dataset {self.dataset_name} in chunks."
|
|
319
|
+
)
|