aidial-rag-eval 0.6.0.dev16__tar.gz → 0.6.0.dev19__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.
Files changed (49) hide show
  1. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/PKG-INFO +27 -5
  2. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/README.md +23 -1
  3. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/pyproject.toml +5 -4
  4. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/dataframe/metrics.py +0 -6
  5. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/evaluate.py +0 -2
  6. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/inference.py +3 -20
  7. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/metric_binds.py +5 -14
  8. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/refusal.py +1 -7
  9. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/utils/segmented_text.py +25 -7
  10. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/LICENSE +0 -0
  11. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/__init__.py +0 -0
  12. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/dataframe/__init__.py +0 -0
  13. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/dataframe/match_facts.py +0 -0
  14. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/dataframe/merge.py +0 -0
  15. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/dataset.py +0 -0
  16. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/facts/__init__.py +0 -0
  17. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/facts/citation.py +0 -0
  18. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/facts/exact.py +0 -0
  19. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/facts/match_str_facts.py +0 -0
  20. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/__init__.py +0 -0
  21. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/__init__.py +0 -0
  22. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/converters/__init__.py +0 -0
  23. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/converters/base_converter.py +0 -0
  24. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/converters/decontextualization_template.py +0 -0
  25. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/converters/llm_decontextualization_converter.py +0 -0
  26. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/inference_scorers/__init__.py +0 -0
  27. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/inference_scorers/base_inference_scorer.py +0 -0
  28. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/inference_scorers/inference_template.py +0 -0
  29. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/inference_scorers/llm_inference_scorer.py +0 -0
  30. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/lambdas.py +0 -0
  31. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/refusal_detectors/__init__.py +0 -0
  32. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/refusal_detectors/base_refusal_detector.py +0 -0
  33. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/refusal_detectors/llm_refusal_detector.py +0 -0
  34. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/refusal_detectors/refusal_template.py +0 -0
  35. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/statement_extractor/__init__.py +0 -0
  36. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/statement_extractor/base_statement_extractor.py +0 -0
  37. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/statement_extractor/llm_statement_extractor.py +0 -0
  38. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/models/statement_extractor/statement_extractor_template.py +0 -0
  39. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/types.py +0 -0
  40. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/utils/__init__.py +0 -0
  41. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/utils/exceptions.py +0 -0
  42. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/generation/utils/progress_bar.py +0 -0
  43. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/metric_binds.py +0 -0
  44. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/metrics.py +0 -0
  45. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/retrieval/__init__.py +0 -0
  46. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/retrieval/metrics.py +0 -0
  47. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/retrieval/types.py +0 -0
  48. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/types.py +0 -0
  49. {aidial_rag_eval-0.6.0.dev16 → aidial_rag_eval-0.6.0.dev19}/src/aidial_rag_eval/utils.py +0 -0
@@ -1,13 +1,13 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: aidial-rag-eval
3
- Version: 0.6.0.dev16
3
+ Version: 0.6.0.dev19
4
4
  Summary: Library for RAG evaluation
5
5
  License: Apache-2.0
6
6
  License-File: LICENSE
7
7
  Keywords: ai
8
8
  Author: EPAM RAIL
9
9
  Author-email: SpecialEPM-DIALDevTeam@epam.com
10
- Requires-Python: >=3.11,<4.0
10
+ Requires-Python: >=3.11,<3.15
11
11
  Classifier: License :: OSI Approved :: Apache Software License
12
12
  Classifier: Programming Language :: Python :: 3
13
13
  Classifier: Programming Language :: Python :: 3.11
@@ -19,12 +19,12 @@ Provides-Extra: s3fs
19
19
  Requires-Dist: fsspec (>=2024.3.1,<2025.0.0)
20
20
  Requires-Dist: langchain-core (>=0.3.80,<2)
21
21
  Requires-Dist: more-itertools (>=10.5.0,<11.0.0)
22
- Requires-Dist: nltk (>=3.8.1,<4.0.0)
23
22
  Requires-Dist: numpy (>=1.26.4,<3)
24
23
  Requires-Dist: pandas (>=2.2.2,<3.0.0)
25
24
  Requires-Dist: pyarrow (>=20.0.0)
26
- Requires-Dist: pydantic (>=2.7.1,<3.0.0)
25
+ Requires-Dist: pydantic (>=2.7.1)
27
26
  Requires-Dist: s3fs (>=2024.3.1,<2025.0.0) ; extra == "s3fs"
27
+ Requires-Dist: spacy (>=3.8.9,<4)
28
28
  Requires-Dist: tqdm (>=4.66.4,<5.0.0)
29
29
  Project-URL: Documentation, https://github.com/epam/ai-dial-rag-eval
30
30
  Project-URL: Homepage, https://github.com/epam/ai-dial-rag-eval
@@ -47,6 +47,27 @@ Install the library using [pip](https://pip.pypa.org/en/stable/getting-started):
47
47
  pip install aidial-rag-eval
48
48
  ```
49
49
 
50
+ ### spaCy language model
51
+
52
+ The generation metrics require the English language model for [spaCy](https://spacy.io/).
53
+ Download it after installation:
54
+
55
+ ```sh
56
+ python -m spacy download en_core_web_sm
57
+ ```
58
+
59
+ Alternatively, you can install the model directly via URL:
60
+
61
+ ```sh
62
+ pip install https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl
63
+ ```
64
+
65
+ Or as a [Poetry](https://python-poetry.org/) dependency:
66
+
67
+ ```toml
68
+ en-core-web-sm = {url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl"}
69
+ ```
70
+
50
71
  ### Example
51
72
 
52
73
  The example of how to get retrieval metrics along with answer inference based on the context.
@@ -113,10 +134,11 @@ In this table:
113
134
 
114
135
  The algorithm is token-intensive. Considering the balance between quality and price, the following models are recommended:
115
136
 
137
+ - **gemini-3.1-flash-lite**
116
138
  - **gemini-2.5-flash-lite**
117
139
  - **gpt-5-mini**
118
- - **gemini-2.0-flash-lite**
119
140
  - **gpt-5-nano**
141
+ - **gpt-5.4-mini**
120
142
 
121
143
  ## Developer environment
122
144
 
@@ -14,6 +14,27 @@ Install the library using [pip](https://pip.pypa.org/en/stable/getting-started):
14
14
  pip install aidial-rag-eval
15
15
  ```
16
16
 
17
+ ### spaCy language model
18
+
19
+ The generation metrics require the English language model for [spaCy](https://spacy.io/).
20
+ Download it after installation:
21
+
22
+ ```sh
23
+ python -m spacy download en_core_web_sm
24
+ ```
25
+
26
+ Alternatively, you can install the model directly via URL:
27
+
28
+ ```sh
29
+ pip install https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl
30
+ ```
31
+
32
+ Or as a [Poetry](https://python-poetry.org/) dependency:
33
+
34
+ ```toml
35
+ en-core-web-sm = {url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl"}
36
+ ```
37
+
17
38
  ### Example
18
39
 
19
40
  The example of how to get retrieval metrics along with answer inference based on the context.
@@ -80,10 +101,11 @@ In this table:
80
101
 
81
102
  The algorithm is token-intensive. Considering the balance between quality and price, the following models are recommended:
82
103
 
104
+ - **gemini-3.1-flash-lite**
83
105
  - **gemini-2.5-flash-lite**
84
106
  - **gpt-5-mini**
85
- - **gemini-2.0-flash-lite**
86
107
  - **gpt-5-nano**
108
+ - **gpt-5.4-mini**
87
109
 
88
110
  ## Developer environment
89
111
 
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "aidial-rag-eval"
3
- version = "0.6.0.dev16"
3
+ version = "0.6.0.dev19"
4
4
  description = "Library for RAG evaluation"
5
5
  authors = ["EPAM RAIL <SpecialEPM-DIALDevTeam@epam.com>"]
6
6
  homepage = "https://github.com/epam/ai-dial-rag-eval"
@@ -13,13 +13,13 @@ repository = "https://github.com/epam/ai-dial-rag-eval"
13
13
 
14
14
 
15
15
  [tool.poetry.dependencies]
16
- python = "^3.11"
16
+ python = ">=3.11, <3.15"
17
17
  numpy = ">=1.26.4,<3"
18
18
  pandas = "^2.2.2"
19
19
  fsspec = "^2024.3.1"
20
20
  s3fs = {version = "^2024.3.1", optional = true}
21
- pydantic = "^2.7.1"
22
- nltk = "^3.8.1"
21
+ pydantic = ">=2.7.1"
22
+ spacy = ">=3.8.9,<4"
23
23
  tqdm = "^4.66.4"
24
24
  langchain-core = ">=0.3.80,<2"
25
25
  more-itertools = "^10.5.0"
@@ -35,6 +35,7 @@ pytest-cov = "^4.1.0"
35
35
  pytest-mock = "^3.14.0"
36
36
  langchain-community = "^0.3.14"
37
37
  langchain-openai = "^0.2.1"
38
+ en-core-web-sm = {url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl"}
38
39
 
39
40
  [tool.poetry.group.lint.dependencies]
40
41
  black = "^24.3.0"
@@ -114,7 +114,6 @@ def calculate_generation_metrics(
114
114
  metric_binds: List[MetricBind],
115
115
  max_concurrency: int = 8,
116
116
  show_progress_bar: bool = True,
117
- auto_download_nltk: bool = True,
118
117
  ) -> pd.DataFrame:
119
118
  """
120
119
  Calculates RAG evaluation generation metrics from df_merged dataframe.
@@ -156,7 +155,6 @@ def calculate_generation_metrics(
156
155
  llm=llm,
157
156
  max_concurrency=max_concurrency,
158
157
  show_progress_bar=show_progress_bar,
159
- auto_download_nltk=auto_download_nltk,
160
158
  ).to_dict(orient="series")
161
159
  )
162
160
  df_metrics = pd.DataFrame(data=metric_results)
@@ -180,7 +178,6 @@ def create_generation_metrics_report(
180
178
  metric_binds: List[MetricBind],
181
179
  max_concurrency: int = 8,
182
180
  show_progress_bar: bool = True,
183
- auto_download_nltk: bool = True,
184
181
  ) -> pd.DataFrame:
185
182
  """
186
183
  Calculates RAG evaluation generation metrics from input dataframes.
@@ -223,7 +220,6 @@ def create_generation_metrics_report(
223
220
  metric_binds,
224
221
  max_concurrency,
225
222
  show_progress_bar,
226
- auto_download_nltk,
227
223
  )
228
224
  return pd.merge(df_merged, df_metrics, left_index=True, right_index=True)
229
225
 
@@ -236,7 +232,6 @@ def create_rag_eval_metrics_report(
236
232
  metric_binds: Optional[List[MetricBind]] = None,
237
233
  max_concurrency: int = 8,
238
234
  show_progress_bar: bool = True,
239
- auto_download_nltk: bool = True,
240
235
  ) -> pd.DataFrame:
241
236
  """
242
237
  Calculates RAG evaluation metrics from input dataframes.
@@ -287,6 +282,5 @@ def create_rag_eval_metrics_report(
287
282
  metric_binds,
288
283
  max_concurrency,
289
284
  show_progress_bar,
290
- auto_download_nltk,
291
285
  )
292
286
  return pd.concat([df_merged, retrieval_metrics, generation_metrics], axis=1)
@@ -21,7 +21,6 @@ def evaluate(
21
21
  metric_binds: Optional[List[MetricBind]] = None,
22
22
  max_concurrency: int = 8,
23
23
  show_progress_bar: bool = True,
24
- auto_download_nltk: bool = False,
25
24
  ) -> Dataset:
26
25
  """
27
26
  Calculates RAG evaluation metrics from input
@@ -83,7 +82,6 @@ def evaluate(
83
82
  metric_binds=metric_binds,
84
83
  max_concurrency=max_concurrency,
85
84
  show_progress_bar=show_progress_bar,
86
- auto_download_nltk=auto_download_nltk,
87
85
  )
88
86
  aggregated_metrics = df_final.mean(numeric_only=True)
89
87
  assert isinstance(aggregated_metrics, pd.Series)
@@ -203,7 +203,6 @@ def _segment_hypotheses(
203
203
  llm: BaseChatModel,
204
204
  max_concurrency: int = 8,
205
205
  show_progress_bar: bool = True,
206
- auto_download_nltk: bool = True,
207
206
  ) -> Tuple[List[SegmentedText], List[Union[SegmentedText, ErrorInfo]]]:
208
207
  """
209
208
  Function that segments hypotheses into hypothesis segments(roughly into
@@ -236,8 +235,7 @@ def _segment_hypotheses(
236
235
  max_concurrency=max_concurrency,
237
236
  )
238
237
  segmented_hypotheses = [
239
- SegmentedText.from_text(text=hypothesis, auto_download_nltk=auto_download_nltk)
240
- for hypothesis in hypotheses
238
+ SegmentedText.from_text(text=hypothesis) for hypothesis in hypotheses
241
239
  ]
242
240
  if show_progress_bar:
243
241
  print("Converting hypothesis...")
@@ -294,14 +292,11 @@ def _extract_statements(
294
292
  def _add_questions_to_premises(
295
293
  premises: List[Premise],
296
294
  questions: Optional[List[Question]],
297
- auto_download_nltk: bool,
298
295
  ) -> List[Premise]:
299
296
  adjusted_premises: List[Premise] = list(premises)
300
297
  if questions is not None:
301
298
  for i, question in enumerate(questions):
302
- question_split = SegmentedText.from_text(
303
- text=question, auto_download_nltk=auto_download_nltk
304
- )
299
+ question_split = SegmentedText.from_text(text=question)
305
300
  adjusted_premises[i] = question_split.segments[-1] + "\n" + premises[i]
306
301
  return adjusted_premises
307
302
 
@@ -314,7 +309,6 @@ def _infer_statements(
314
309
  list_documents: Optional[List[Documents]] = None,
315
310
  max_concurrency: int = 8,
316
311
  show_progress_bar: bool = True,
317
- auto_download_nltk: bool = True,
318
312
  ) -> List[List[Tuple[InferenceInputs, InferenceScore]]]:
319
313
  """
320
314
  Function that infers statements.
@@ -358,9 +352,7 @@ def _infer_statements(
358
352
  For items that failed in any stage, InferenceScore will have inference=None
359
353
  and the error field set.
360
354
  """
361
- adjusted_premises = _add_questions_to_premises(
362
- premises, questions, auto_download_nltk
363
- )
355
+ adjusted_premises = _add_questions_to_premises(premises, questions)
364
356
 
365
357
  document_names: List[JoinedDocumentsName] = (
366
358
  [""] * len(statements)
@@ -391,14 +383,12 @@ def segment_hypotheses(
391
383
  llm: BaseChatModel,
392
384
  max_concurrency: int = 8,
393
385
  show_progress_bar: bool = True,
394
- auto_download_nltk: bool = True,
395
386
  ) -> List[Union[SegmentedText, ErrorInfo]]:
396
387
  return _segment_hypotheses(
397
388
  hypotheses=hypotheses,
398
389
  llm=llm,
399
390
  max_concurrency=max_concurrency,
400
391
  show_progress_bar=show_progress_bar,
401
- auto_download_nltk=auto_download_nltk,
402
392
  )[1]
403
393
 
404
394
 
@@ -426,7 +416,6 @@ def infer_statements(
426
416
  list_documents: Optional[List[Documents]] = None,
427
417
  max_concurrency: int = 8,
428
418
  show_progress_bar: bool = True,
429
- auto_download_nltk: bool = True,
430
419
  ) -> List[List[Tuple[InferenceInputs, InferenceScore]]]:
431
420
  return _infer_statements(
432
421
  premises=premises,
@@ -436,7 +425,6 @@ def infer_statements(
436
425
  list_documents=list_documents,
437
426
  max_concurrency=max_concurrency,
438
427
  show_progress_bar=show_progress_bar,
439
- auto_download_nltk=auto_download_nltk,
440
428
  )
441
429
 
442
430
 
@@ -484,7 +472,6 @@ def calculate_batch_inference(
484
472
  list_documents: Optional[List[Documents]] = None,
485
473
  max_concurrency: int = 8,
486
474
  show_progress_bar: bool = True,
487
- auto_download_nltk: bool = True,
488
475
  ) -> List[InferenceReturn]:
489
476
  """
490
477
  Calculates pairwise the inference of a hypotheses from a premises.
@@ -530,7 +517,6 @@ def calculate_batch_inference(
530
517
  llm=llm,
531
518
  max_concurrency=max_concurrency,
532
519
  show_progress_bar=show_progress_bar,
533
- auto_download_nltk=auto_download_nltk,
534
520
  )
535
521
  statements: List[Union[HypothesisStatements, ErrorInfo]] = _extract_statements(
536
522
  segmented_hypotheses=decontextualized_segmented_hypotheses,
@@ -547,7 +533,6 @@ def calculate_batch_inference(
547
533
  list_documents=list_documents,
548
534
  max_concurrency=max_concurrency,
549
535
  show_progress_bar=show_progress_bar,
550
- auto_download_nltk=auto_download_nltk,
551
536
  )
552
537
  )
553
538
 
@@ -585,7 +570,6 @@ def calculate_inference(
585
570
  documents: Optional[Documents] = None,
586
571
  max_concurrency: int = 8,
587
572
  show_progress_bar: bool = True,
588
- auto_download_nltk: bool = True,
589
573
  ) -> InferenceReturn:
590
574
  """
591
575
  Calculates the inference of a hypothesis from a premise.
@@ -631,6 +615,5 @@ def calculate_inference(
631
615
  list_documents=list_documents,
632
616
  max_concurrency=max_concurrency,
633
617
  show_progress_bar=show_progress_bar,
634
- auto_download_nltk=auto_download_nltk,
635
618
  )
636
619
  return inference_returns[0]
@@ -33,7 +33,6 @@ def _wrapped_dataframe_inference(
33
33
  document_column: Optional[str] = None,
34
34
  max_concurrency: int = 8,
35
35
  show_progress_bar: bool = True,
36
- auto_download_nltk: bool = True,
37
36
  ) -> pd.DataFrame:
38
37
  inference_returns = calculate_batch_inference(
39
38
  premises=_get_column_as_list_str(df_merged, premise_column),
@@ -51,7 +50,6 @@ def _wrapped_dataframe_inference(
51
50
  ),
52
51
  max_concurrency=max_concurrency,
53
52
  show_progress_bar=show_progress_bar,
54
- auto_download_nltk=auto_download_nltk,
55
53
  )
56
54
  return pd.DataFrame(
57
55
  [dataclasses.asdict(inference_return) for inference_return in inference_returns]
@@ -65,14 +63,12 @@ def _wrapped_dataframe_refusal(
65
63
  prefix: str,
66
64
  max_concurrency: int = 8,
67
65
  show_progress_bar: bool = True,
68
- auto_download_nltk: bool = True,
69
66
  ) -> pd.DataFrame:
70
67
  refusal_returns = calculate_batch_refusal(
71
68
  answers=_get_column_as_list_str(df_merged, answer_column),
72
69
  llm=llm,
73
70
  max_concurrency=max_concurrency,
74
71
  show_progress_bar=show_progress_bar,
75
- auto_download_nltk=auto_download_nltk,
76
72
  )
77
73
  return pd.DataFrame(
78
74
  [dataclasses.asdict(refusal) for refusal in refusal_returns]
@@ -80,7 +76,7 @@ def _wrapped_dataframe_refusal(
80
76
 
81
77
 
82
78
  def context_to_answer_inference(
83
- df_merged, llm, max_concurrency, show_progress_bar, auto_download_nltk, **kwargs
79
+ df_merged, llm, max_concurrency, show_progress_bar, **kwargs
84
80
  ) -> pd.DataFrame:
85
81
  return _wrapped_dataframe_inference(
86
82
  df_merged=df_merged,
@@ -98,12 +94,11 @@ def context_to_answer_inference(
98
94
  document_column=MergedColumns.DOCUMENTS,
99
95
  max_concurrency=max_concurrency,
100
96
  show_progress_bar=show_progress_bar,
101
- auto_download_nltk=auto_download_nltk,
102
97
  )
103
98
 
104
99
 
105
100
  def answer_to_ground_truth_inference(
106
- df_merged, llm, max_concurrency, show_progress_bar, auto_download_nltk, **kwargs
101
+ df_merged, llm, max_concurrency, show_progress_bar, **kwargs
107
102
  ) -> pd.DataFrame:
108
103
  return _wrapped_dataframe_inference(
109
104
  df_merged=df_merged,
@@ -115,12 +110,11 @@ def answer_to_ground_truth_inference(
115
110
  document_column=MergedColumns.DOCUMENTS,
116
111
  max_concurrency=max_concurrency,
117
112
  show_progress_bar=show_progress_bar,
118
- auto_download_nltk=auto_download_nltk,
119
113
  )
120
114
 
121
115
 
122
116
  def ground_truth_to_answer_inference(
123
- df_merged, llm, max_concurrency, show_progress_bar, auto_download_nltk, **kwargs
117
+ df_merged, llm, max_concurrency, show_progress_bar, **kwargs
124
118
  ) -> pd.DataFrame:
125
119
  return _wrapped_dataframe_inference(
126
120
  df_merged=df_merged,
@@ -132,12 +126,11 @@ def ground_truth_to_answer_inference(
132
126
  document_column=MergedColumns.DOCUMENTS,
133
127
  max_concurrency=max_concurrency,
134
128
  show_progress_bar=show_progress_bar,
135
- auto_download_nltk=auto_download_nltk,
136
129
  )
137
130
 
138
131
 
139
132
  def answer_refusal(
140
- df_merged, llm, max_concurrency, show_progress_bar, auto_download_nltk, **kwargs
133
+ df_merged, llm, max_concurrency, show_progress_bar, **kwargs
141
134
  ) -> pd.DataFrame:
142
135
  return _wrapped_dataframe_refusal(
143
136
  df_merged=df_merged,
@@ -146,12 +139,11 @@ def answer_refusal(
146
139
  prefix=ANSWER_REFUSAL_PREFIX,
147
140
  max_concurrency=max_concurrency,
148
141
  show_progress_bar=show_progress_bar,
149
- auto_download_nltk=auto_download_nltk,
150
142
  )
151
143
 
152
144
 
153
145
  def ground_truth_refusal(
154
- df_merged, llm, max_concurrency, show_progress_bar, auto_download_nltk, **kwargs
146
+ df_merged, llm, max_concurrency, show_progress_bar, **kwargs
155
147
  ) -> pd.DataFrame:
156
148
  return _wrapped_dataframe_refusal(
157
149
  df_merged=df_merged,
@@ -160,7 +152,6 @@ def ground_truth_refusal(
160
152
  prefix=GT_ANSWER_REFUSAL_PREFIX,
161
153
  max_concurrency=max_concurrency,
162
154
  show_progress_bar=show_progress_bar,
163
- auto_download_nltk=auto_download_nltk,
164
155
  )
165
156
 
166
157
 
@@ -15,7 +15,6 @@ def calculate_batch_refusal(
15
15
  llm: BaseChatModel,
16
16
  max_concurrency: int = 8,
17
17
  show_progress_bar: bool = True,
18
- auto_download_nltk: bool = True,
19
18
  ) -> List[RefusalReturn]:
20
19
  """
21
20
  Checks if the answers are answer refusal.
@@ -41,10 +40,7 @@ def calculate_batch_refusal(
41
40
  Returns the list of the answer refusals.
42
41
  """
43
42
  detector = LLMRefusalDetector(llm, max_concurrency)
44
- answers_split = [
45
- SegmentedText.from_text(text=answer, auto_download_nltk=auto_download_nltk)
46
- for answer in answers
47
- ]
43
+ answers_split = [SegmentedText.from_text(text=answer) for answer in answers]
48
44
  # As a heuristic, we send only the first 3 segments in the prompt.
49
45
  # We believe that if there are 3 whole segments with information
50
46
  # that is not related to refusal to answer,
@@ -64,7 +60,6 @@ def calculate_refusal(
64
60
  llm: BaseChatModel,
65
61
  max_concurrency: int = 8,
66
62
  show_progress_bar: bool = True,
67
- auto_download_nltk: bool = True,
68
63
  ) -> RefusalReturn:
69
64
  """
70
65
  Checks if the answer is answer refusal.
@@ -94,6 +89,5 @@ def calculate_refusal(
94
89
  llm=llm,
95
90
  max_concurrency=max_concurrency,
96
91
  show_progress_bar=show_progress_bar,
97
- auto_download_nltk=auto_download_nltk,
98
92
  )
99
93
  return refusal_returns[0]
@@ -1,11 +1,32 @@
1
1
  from typing import Callable, List, Optional, Tuple
2
2
 
3
- import nltk
4
- from more_itertools import chunked
3
+ import spacy
4
+ from more_itertools import chunked, constrained_batches
5
5
 
6
6
  from aidial_rag_eval.generation.types import Text, TextSegment
7
7
 
8
8
  Splitter = Callable[[TextSegment], List[TextSegment]]
9
+
10
+ _nlp = spacy.load("en_core_web_sm")
11
+
12
+ _simple_nlp = spacy.blank("en")
13
+ _simple_nlp.add_pipe("sentencizer")
14
+ _simple_nlp.max_length = 1_000_000_000
15
+
16
+
17
+ def _sent_tokenize(text: str) -> List[str]:
18
+ max_len = _nlp.max_length
19
+ sents = list(_simple_nlp(text).sents)
20
+
21
+ chunks = [
22
+ "".join(s.text_with_ws for s in batch)
23
+ for batch in constrained_batches(sents, max_len, get_len=lambda s: len(s.text))
24
+ ]
25
+ return [
26
+ s.text.strip() for chunk in chunks for s in _nlp(chunk).sents if s.text.strip()
27
+ ]
28
+
29
+
9
30
  SegmentChecker = Callable[[TextSegment], bool]
10
31
 
11
32
 
@@ -80,10 +101,7 @@ class SegmentedText:
80
101
  self.delimiters = delimiters.copy()
81
102
 
82
103
  @classmethod
83
- def from_text(cls, text: Text, auto_download_nltk: bool = True) -> "SegmentedText":
84
- if auto_download_nltk:
85
- nltk.download("punkt_tab", quiet=True)
86
-
104
+ def from_text(cls, text: Text) -> "SegmentedText":
87
105
  max_len = 500
88
106
  min_len = 10
89
107
  conditional_splitters: List[Splitter] = [
@@ -93,7 +111,7 @@ class SegmentedText:
93
111
  ]
94
112
  segmented_text = cls([text], [])
95
113
  segmented_text = apply_splitter_to_segmented_text(
96
- segmented_text, nltk.sent_tokenize
114
+ segmented_text, _sent_tokenize
97
115
  )
98
116
  for splitter in conditional_splitters:
99
117
  segmented_text = apply_splitter_to_segmented_text(