Merge pull request #389 from ggozad/chore/clean-up-evaluations
Clean up and update evaluations.
This commit is contained in:
commit
d5e5733f67
26 changed files with 362 additions and 796 deletions
|
|
@ -36,13 +36,6 @@ Active datasets:
|
|||
| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB |
|
||||
| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB |
|
||||
|
||||
Inactive (kept downloadable, not currently maintained):
|
||||
|
||||
| Dataset | Size |
|
||||
|---------|------|
|
||||
| `repliqa` | ~30MB |
|
||||
| `hotpotqa` | ~331MB |
|
||||
|
||||
After downloading, run benchmarks with `--skip-db` to use the pre-built database:
|
||||
|
||||
```bash
|
||||
|
|
@ -85,20 +78,13 @@ evaluations:
|
|||
|
||||
### Retrieval Metrics
|
||||
|
||||
**Mean Reciprocal Rank (MRR)** - Used when each query has exactly one relevant document.
|
||||
|
||||
- For each query, find the rank (position) of the first relevant document in top-K results
|
||||
- Reciprocal rank = `1/rank` (e.g., rank 3 → 1/3 ≈ 0.333)
|
||||
- If not found in top-K, score is 0
|
||||
- MRR is the mean across all queries
|
||||
- Range: 0 (never found) to 1 (always at rank 1)
|
||||
|
||||
**Mean Average Precision (MAP)** - Used when queries have multiple relevant documents.
|
||||
**Mean Average Precision (MAP)** scores ranked retrieval results against the gold `expected_uris`.
|
||||
|
||||
- For each relevant document at position k, calculate precision@k = (relevant docs in top k) / k
|
||||
- Average Precision (AP) = mean of these precision values / total relevant documents
|
||||
- Average Precision (AP) = sum of these precision values / total relevant documents
|
||||
- MAP is the mean of AP scores across all queries
|
||||
- Range: 0 to 1. Rewards ranking relevant documents higher
|
||||
- For single-doc queries this collapses to `1/rank` (i.e. reciprocal rank)
|
||||
|
||||
### QA Accuracy
|
||||
|
||||
|
|
@ -108,7 +94,7 @@ We picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibrat
|
|||
|
||||
### Citation Retrieval
|
||||
|
||||
Alongside QA accuracy, a second metric scores the URIs the skill registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MRR / MAP math as raw retrieval. The score key is `cited_mrr` for single-doc datasets and `cited_map` for multi-doc. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
|
||||
Alongside QA accuracy, a second metric scores the URIs the skill registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
|
||||
|
||||
This is computed alongside QA accuracy from the same skill run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the skill grounded its answer on it.
|
||||
|
||||
|
|
@ -116,18 +102,6 @@ This is computed alongside QA accuracy from the same skill run, no extra invocat
|
|||
|
||||
Numbers measured under the current pinned judge (`ollama:qwen3.6`) on a recent `haiku.rag` version.
|
||||
|
||||
### Wix
|
||||
|
||||
[WixQA](https://huggingface.co/datasets/Wix/WixQA) is real customer support questions paired with curated answers. 200 cases.
|
||||
|
||||
`evaluations run wix --target rag-skill` runs the RAG skill end-to-end and produces both QA accuracy and a citation retrieval metric (`cited_map`) computed from the URIs the skill registered via the `cite` tool against the gold `expected_uris`.
|
||||
|
||||
| Skill model | Reranker | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------|------------------------|-------------|------------------|
|
||||
| `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | 0.87 | 0.38 |
|
||||
|
||||
*Measured on haiku.rag v0.48.0 with `qwen3-embedding:4b` (vLLM, dim 2560), `chunk_size=256`, `search.limit=5`. Judged by `vllm:Qwen3.6-35B-A3B-NVFP4` (qwen3.6 family, NVFP4 quant served via vLLM rather than the default Ollama). 172 / 198 completed cases (2 errored).*
|
||||
|
||||
### OpenRAG Bench (ORB)
|
||||
|
||||
[OpenRAG Bench](https://huggingface.co/datasets/vectara/open_ragbench) contains ArXiv research papers with multimodal question-answering pairs. Queries include both text-based and image-based questions, testing retrieval and reasoning over visual content like figures, charts, and diagrams. Each query maps to one relevant document.
|
||||
|
|
@ -141,114 +115,49 @@ Two approaches are benchmarked separately:
|
|||
|
||||
##### Retrieval (MAP)
|
||||
|
||||
| Embedding Model | Source bucket | Cases | MAP |
|
||||
|------------------------------|--------------------|------:|-------:|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | text only | 1914 | 0.9801 |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | text + image | 763 | 0.9720 |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | text + table | 148 | 0.9786 |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | text + table+image | 220 | 0.9720 |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | **all** | 3045 | **0.9774** |
|
||||
| Embedding Model | Cases | MAP |
|
||||
|------------------------------------------|------:|-------:|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | 3045 | 0.9774 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | 3045 | 0.9709 |
|
||||
|
||||
##### QA Accuracy
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
| Embedding Model | Skill model | Reranker | Source bucket | Cases | Accuracy |
|
||||
|------------------------------|-----------------------------------|------------------------|---------------|------:|---------:|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `ollama:qwen3.6` (vision) | none | text only | 682 | 96.9 % |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `ollama:qwen3.6` (vision) | none | with image | 299 | 91.3 % |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | text | 894 | 88.5 % |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | text+image | 341 | 88.0 % |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | text+table | 72 | 88.9 % |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | text+table+image | 102 | 93.1 % |
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | **all** | 1409 | **88.7 %** |
|
||||
| Embedding Model | Target | Skill model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------------------|-----------------|------------------------------|------:|-------------|------------------|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `rag-skill` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-skill` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.93 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-skill`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.94 | 0.35 |
|
||||
|
||||
*Measured on haiku.rag v0.50.0, no reranker (no multimodal reranker available), judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Qwen3-VL covered 1409 / 3045 cases and predates the `cited_map` evaluator. The analysis skill answers via Python sandbox over the document VFS and only cites in ~35 % of cases, which depresses `cited_map`.*
|
||||
|
||||
#### Text embedder + VLM picture descriptions
|
||||
|
||||
##### Retrieval (MAP)
|
||||
|
||||
| Embedding Model | VLM | Source bucket | Cases | MAP |
|
||||
|----------------------|----------------------|---------------|------:|-----------:|
|
||||
| `qwen3-embedding:4b` | Ollama / ministral-3 | **all** | 3045 | **0.9722** |
|
||||
| Embedding Model | VLM | Reranker | Cases | MAP |
|
||||
|------------------------------------------|----------------------|------------------------|------:|-------:|
|
||||
| `qwen3-embedding:4b` | Ollama / ministral-3 | `mxbai-rerank-base-v2` | 3045 | 0.9834 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `mxbai-rerank-base-v2` | 3045 | 0.9863 |
|
||||
|
||||
*Measured on haiku.rag v0.45.0.*
|
||||
*Measured on haiku.rag v0.50.0.*
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
| Embedding Model | VLM | Skill model | QA accuracy | Mean `cited_map` |
|
||||
|------------------------|----------------------|------------------------------|-------------|------------------|
|
||||
| `qwen3-embedding:4b` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 0.88 | 0.89 |
|
||||
| Embedding Model | VLM | Skill model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------------------|----------------------|------------------------------|------:|-------------|------------------|
|
||||
| `qwen3-embedding:4b` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.80 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 2836 | 0.96 | 0.81 |
|
||||
|
||||
*Measured on haiku.rag v0.48.0, with `mxbai-rerank-base-v2`, on all 3045 cases. Judged by `vllm:Qwen3.6-35B-A3B-NVFP4`.*
|
||||
*Measured on haiku.rag v0.50.0 with `mxbai-rerank-base-v2`, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Nemotron ran 2836 / 3045 cases (run stopped early; numbers stable from ~12% onward).*
|
||||
|
||||
## Inactive datasets
|
||||
### Wix
|
||||
|
||||
The benchmarks below are not currently maintained. Numbers were measured against earlier `haiku.rag` versions and an older pinned judge (`ollama:gpt-oss`), before the skill workflow became the only path. Retrieval tables don't depend on the judge, but the QA tables aren't reproducible against the current skill-only setup. We may revive them.
|
||||
[WixQA](https://huggingface.co/datasets/Wix/WixQA) is real customer support questions paired with curated answers. 200 cases.
|
||||
|
||||
### RepliQA
|
||||
`evaluations run wix --target rag-skill` runs the RAG skill end-to-end and produces both QA accuracy and a citation retrieval metric (`cited_map`) computed from the URIs the skill registered via the `cite` tool against the gold `expected_uris`.
|
||||
|
||||
[RepliQA](https://huggingface.co/datasets/ServiceNow/repliqa) contains synthetic news stories with question-answer pairs. We use `News Stories` from `repliqa_3` (1035 documents). Each question has exactly one relevant document, so we use MRR for retrieval evaluation.
|
||||
| Skill model | Reranker | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------|------------------------|-------------|------------------|
|
||||
| `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | 0.87 | 0.38 |
|
||||
|
||||
#### Retrieval (MRR)
|
||||
|
||||
| Embedding Model | MRR | Reranker |
|
||||
|-------------------------------|------|----------|
|
||||
| Ollama / `qwen3-embedding:8b` | 0.91 | - |
|
||||
|
||||
*Measured on haiku.rag v0.19.6.*
|
||||
|
||||
#### QA Accuracy
|
||||
|
||||
| Embedding Model | QA Model | Accuracy | Reranker |
|
||||
|------------------------------|----------------------------------|----------|------------------------|
|
||||
| Ollama / `qwen3-embedding:4b` | Ollama / `gpt-oss` - no thinking | 0.82 | None |
|
||||
| Ollama / `qwen3-embedding:8b` | Ollama / `gpt-oss` - thinking | 0.89 | None |
|
||||
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` - thinking | 0.85 | None |
|
||||
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` - thinking | 0.87 | `mxbai-rerank-base-v2` |
|
||||
| Ollama / `mxbai-embed-large` | Ollama / `qwen3:0.6b` | 0.28 | None |
|
||||
|
||||
*Measured on haiku.rag v0.19.6, judged by `ollama:gpt-oss`.*
|
||||
|
||||
Note the significant degradation when very small models are used such as `qwen3:0.6b`.
|
||||
|
||||
### HotpotQA
|
||||
|
||||
[HotpotQA](https://huggingface.co/datasets/hotpotqa/hotpot_qa) is a multi-hop question answering dataset requiring reasoning over multiple Wikipedia paragraphs. Each question requires evidence from 2+ documents, making it ideal for testing retrieval and reasoning capabilities. We use MAP for retrieval evaluation since queries have multiple relevant documents.
|
||||
|
||||
#### Retrieval (MAP)
|
||||
|
||||
| Embedding Model | MAP | Reranker |
|
||||
|----------------------|------|----------|
|
||||
| `qwen3-embedding:4b` | 0.69 | none |
|
||||
|
||||
*Measured on haiku.rag v0.20.2.*
|
||||
|
||||
#### QA Accuracy
|
||||
|
||||
| Embedding Model | QA Model | Accuracy |
|
||||
|----------------------|--------------------------|----------|
|
||||
| `qwen3-embedding:4b` | `gpt-oss:20b` - thinking | 0.86 |
|
||||
|
||||
*Measured on haiku.rag v0.20.2, judged by `ollama:gpt-oss`.*
|
||||
|
||||
### Wix (historical, plain text and HTML)
|
||||
|
||||
Earlier Wix runs measured under different chunk settings and reranker combinations, against the older `gpt-oss` judge.
|
||||
|
||||
#### Retrieval (MAP)
|
||||
|
||||
| Embedding Model | Chunk size | MAP | Reranker | Notes |
|
||||
|------------------------|------------|------|------------------------|------------------------------|
|
||||
| `qwen3-embedding:4b` | 256 | 0.34 | None | html, `chunk-radius=2` |
|
||||
| `qwen3-embedding:4b` | 256 | 0.39 | `mxbai-rerank-base-v2` | html, `chunk-radius=2` |
|
||||
| `qwen3-embedding:4b` | 256 | 0.43 | None | plain text, `chunk-radius=0` |
|
||||
| `qwen3-embedding:4b` | 512 | 0.45 | None | plain text, `chunk-radius=0` |
|
||||
|
||||
*Measured on haiku.rag v0.27.2.*
|
||||
|
||||
#### QA Accuracy
|
||||
|
||||
| Embedding Model | Chunk size | QA Model | Accuracy | Notes |
|
||||
|----------------------|------------|-----------------------------|----------|------------------------------|
|
||||
| `qwen3-embedding:4b` | 256 | `gpt-oss:20b` - no thinking | 0.80 | html, `chunk-radius=2` |
|
||||
| `qwen3-embedding:4b` | 256 | `gpt-oss:20b` - no thinking | 0.83 | html, `chunk-radius=2`, `jinaai/jina-reranker-v3` |
|
||||
|
||||
*Measured on haiku.rag v0.27.2, judged by `ollama:gpt-oss`.*
|
||||
*Measured on haiku.rag v0.48.0 with `qwen3-embedding:4b` (vLLM, dim 2560), `chunk_size=256`, `search.limit=5`. Judged by `vllm:Qwen3.6-35B-A3B-NVFP4` (qwen3.6 family, NVFP4 quant served via vLLM rather than the default Ollama). 172 / 198 completed cases (2 errored).*
|
||||
|
|
|
|||
|
|
@ -8,9 +8,7 @@ This package is not published to PyPI and is only used for development and testi
|
|||
|
||||
Contains evaluation scripts for benchmarking RAG retrieval and QA performance. Available datasets:
|
||||
|
||||
- RepliQA (`repliqa`)
|
||||
- WiX (`wix`)
|
||||
- HotpotQA (`hotpotqa`)
|
||||
- OpenRAG Bench, two variants:
|
||||
- `orb_text` — text embedder (`qwen3-embedding:4b`, 2560-dim) with VLM picture descriptions baked into chunk content at ingest. Use for text-only retrieval/QA against figure-rich corpora.
|
||||
- `orb_multimodal` — multimodal embedder (`qwen3-vl-embedding-8b`, 4096-dim) with picture vectors in the same space as text. Use for cross-modal retrieval (text-as-query → figure hits, image-as-query) and vision QA where the figure itself is the answer.
|
||||
|
|
@ -21,24 +19,24 @@ After installing the package, you can run evaluations using the `evaluations` co
|
|||
|
||||
```bash
|
||||
# Run retrieval + QA benchmarks
|
||||
evaluations run repliqa
|
||||
evaluations run wix
|
||||
evaluations run orb_text
|
||||
|
||||
# Use a custom config file
|
||||
evaluations run repliqa --config /path/to/haiku.rag.yaml
|
||||
evaluations run wix --config /path/to/haiku.rag.yaml
|
||||
|
||||
# Override the database path
|
||||
evaluations run repliqa --db /path/to/custom.lancedb
|
||||
evaluations run wix --db /path/to/custom.lancedb
|
||||
|
||||
# Skip database population and run only benchmarks
|
||||
evaluations run repliqa --skip-db
|
||||
evaluations run wix --skip-db
|
||||
|
||||
# Skip specific benchmarks
|
||||
evaluations run repliqa --skip-retrieval
|
||||
evaluations run repliqa --skip-qa
|
||||
evaluations run wix --skip-retrieval
|
||||
evaluations run wix --skip-qa
|
||||
|
||||
# Limit the number of test cases
|
||||
evaluations run repliqa --limit 100
|
||||
evaluations run wix --limit 100
|
||||
```
|
||||
|
||||
### Choosing the target
|
||||
|
|
@ -54,24 +52,23 @@ evaluations run wix --target analysis-skill --skill-model ollama:gpt-oss
|
|||
|
||||
`--skill-model "provider:name"` overrides the skill model independently from
|
||||
the judge (defaults to `qa.model`, or `analysis.model` when set for the
|
||||
analysis-skill target). A citation retrieval metric (`cited_mrr` / `cited_map`)
|
||||
is computed alongside QA accuracy from the URIs the skill registered via the
|
||||
`cite` tool.
|
||||
analysis-skill target). A citation retrieval metric (`cited_map`) is computed
|
||||
alongside QA accuracy from the URIs the skill registered via the `cite` tool.
|
||||
|
||||
### Pre-built Databases
|
||||
|
||||
Download pre-built evaluation databases from HuggingFace:
|
||||
|
||||
```bash
|
||||
evaluations download repliqa
|
||||
evaluations download wix
|
||||
evaluations download all
|
||||
evaluations download repliqa --force
|
||||
evaluations download wix --force
|
||||
```
|
||||
|
||||
Upload databases (maintainer only):
|
||||
|
||||
```bash
|
||||
evaluations upload repliqa
|
||||
evaluations upload wix
|
||||
evaluations upload all
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -19,9 +19,7 @@ from evaluations.datasets import DATASETS
|
|||
from evaluations.evaluators import (
|
||||
ANSWER_EQUIVALENCE_RUBRIC,
|
||||
CitationMAPEvaluator,
|
||||
CitationMRREvaluator,
|
||||
MAPEvaluator,
|
||||
MRREvaluator,
|
||||
)
|
||||
from evaluations.skill_runner import SkillFactory, run_skill_question
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
|
@ -30,11 +28,6 @@ from haiku.rag.config.models import ModelConfig
|
|||
from haiku.rag.logging import configure_cli_logging
|
||||
from haiku.rag.utils import get_model, parse_model_option
|
||||
|
||||
_CITATION_EVALUATORS: dict[type[Evaluator], type[Evaluator]] = {
|
||||
MRREvaluator: CitationMRREvaluator,
|
||||
MAPEvaluator: CitationMAPEvaluator,
|
||||
}
|
||||
|
||||
Target = Literal["rag-skill", "analysis-skill"]
|
||||
TARGETS: tuple[Target, ...] = ("rag-skill", "analysis-skill")
|
||||
|
||||
|
|
@ -306,10 +299,9 @@ def _skill_factory_for_target(target: Target) -> SkillFactory:
|
|||
|
||||
def _citation_evaluator_for(retrieval_evaluator: Evaluator | None) -> Evaluator | None:
|
||||
"""Return the citation-scoring twin of the dataset's retrieval evaluator."""
|
||||
if retrieval_evaluator is None:
|
||||
return None
|
||||
twin = _CITATION_EVALUATORS.get(type(retrieval_evaluator))
|
||||
return twin() if twin is not None else None
|
||||
if isinstance(retrieval_evaluator, MAPEvaluator):
|
||||
return CitationMAPEvaluator()
|
||||
return None
|
||||
|
||||
|
||||
def _attach_relevant_uris(
|
||||
|
|
|
|||
|
|
@ -1,16 +1,12 @@
|
|||
from evaluations.config import DatasetSpec
|
||||
|
||||
from .hotpotqa import HOTPOTQA_SPEC
|
||||
from .open_rag_bench import ORB_MULTIMODAL_SPEC, ORB_TEXT_SPEC
|
||||
from .repliqa import REPLIQA_SPEC
|
||||
from .wix import WIX_SPEC
|
||||
|
||||
DATASETS: dict[str, DatasetSpec] = {
|
||||
spec.key: spec
|
||||
for spec in (
|
||||
REPLIQA_SPEC,
|
||||
WIX_SPEC,
|
||||
HOTPOTQA_SPEC,
|
||||
ORB_TEXT_SPEC,
|
||||
ORB_MULTIMODAL_SPEC,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,108 +0,0 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, cast
|
||||
|
||||
from datasets import Dataset, load_dataset
|
||||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
|
||||
from evaluations.evaluators import MAPEvaluator
|
||||
|
||||
|
||||
def load_hotpotqa_validation() -> Dataset:
|
||||
dataset_dict = load_dataset("hotpotqa/hotpot_qa", "distractor")
|
||||
return dataset_dict["validation"]
|
||||
|
||||
|
||||
def extract_unique_documents(dataset: Dataset) -> list[dict[str, Any]]:
|
||||
"""Extract unique documents from all context paragraphs, deduplicated by title."""
|
||||
seen_titles: set[str] = set()
|
||||
documents: list[dict[str, Any]] = []
|
||||
|
||||
for sample in dataset:
|
||||
sample = cast(Mapping[str, Any], sample)
|
||||
context = sample["context"]
|
||||
titles = context["title"]
|
||||
sentences_list = context["sentences"]
|
||||
|
||||
for title, sentences in zip(titles, sentences_list):
|
||||
if title in seen_titles:
|
||||
continue
|
||||
seen_titles.add(title)
|
||||
content = " ".join(sentences)
|
||||
documents.append({"title": title, "content": content})
|
||||
|
||||
return documents
|
||||
|
||||
|
||||
_cached_documents: list[dict[str, Any]] | None = None
|
||||
|
||||
|
||||
def load_hotpotqa_documents() -> list[dict[str, Any]]:
|
||||
"""Load and cache unique documents from HotpotQA."""
|
||||
global _cached_documents
|
||||
if _cached_documents is None:
|
||||
dataset = load_hotpotqa_validation()
|
||||
_cached_documents = extract_unique_documents(dataset)
|
||||
return _cached_documents
|
||||
|
||||
|
||||
def document_loader() -> Dataset:
|
||||
"""Return documents as a Dataset-like iterable."""
|
||||
docs = load_hotpotqa_documents()
|
||||
return Dataset.from_list(docs)
|
||||
|
||||
|
||||
def map_hotpotqa_document(doc: Mapping[str, Any]) -> DocumentPayload:
|
||||
return DocumentPayload(
|
||||
uri=doc["title"],
|
||||
content=doc["content"],
|
||||
title=doc["title"],
|
||||
)
|
||||
|
||||
|
||||
def map_hotpotqa_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
|
||||
supporting_facts = doc["supporting_facts"]
|
||||
titles = supporting_facts["title"]
|
||||
if not titles:
|
||||
return None
|
||||
|
||||
unique_titles = tuple(dict.fromkeys(titles))
|
||||
return RetrievalSample(
|
||||
question=doc["question"],
|
||||
expected_uris=unique_titles,
|
||||
)
|
||||
|
||||
|
||||
def build_hotpotqa_case(
|
||||
index: int, doc: Mapping[str, Any]
|
||||
) -> Case[str, str, dict[str, str]]:
|
||||
question_id = doc["id"]
|
||||
question_type = doc["type"]
|
||||
level = doc["level"]
|
||||
|
||||
case_name = f"{index}_{question_id}"
|
||||
|
||||
return Case(
|
||||
name=case_name,
|
||||
inputs=doc["question"],
|
||||
expected_output=doc["answer"],
|
||||
metadata={
|
||||
"question_id": str(question_id),
|
||||
"type": str(question_type),
|
||||
"level": str(level),
|
||||
"case_index": str(index),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
HOTPOTQA_SPEC = DatasetSpec(
|
||||
key="hotpotqa",
|
||||
db_filename="hotpotqa.lancedb",
|
||||
document_loader=document_loader,
|
||||
document_mapper=map_hotpotqa_document,
|
||||
qa_loader=load_hotpotqa_validation,
|
||||
qa_case_builder=build_hotpotqa_case,
|
||||
retrieval_loader=load_hotpotqa_validation,
|
||||
retrieval_mapper=map_hotpotqa_retrieval,
|
||||
retrieval_evaluator=MAPEvaluator(),
|
||||
)
|
||||
|
|
@ -1,60 +0,0 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
from datasets import Dataset, load_dataset
|
||||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
|
||||
from evaluations.evaluators import MRREvaluator
|
||||
|
||||
|
||||
def load_repliqa_corpus() -> Dataset:
|
||||
dataset_dict = load_dataset("ServiceNow/repliqa")
|
||||
dataset = dataset_dict["repliqa_3"]
|
||||
return dataset.filter(lambda doc: doc["document_topic"] == "News Stories")
|
||||
|
||||
|
||||
def map_repliqa_document(doc: Mapping[str, Any]) -> DocumentPayload:
|
||||
return DocumentPayload(
|
||||
uri=str(doc["document_id"]),
|
||||
content=doc["document_extracted"],
|
||||
)
|
||||
|
||||
|
||||
def map_repliqa_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
|
||||
expected_answer = doc["answer"]
|
||||
if expected_answer == "The answer is not found in the document.":
|
||||
return None
|
||||
return RetrievalSample(
|
||||
question=doc["question"],
|
||||
expected_uris=(str(doc["document_id"]),),
|
||||
)
|
||||
|
||||
|
||||
def build_repliqa_case(
|
||||
index: int, doc: Mapping[str, Any]
|
||||
) -> Case[str, str, dict[str, str]]:
|
||||
document_id = doc["document_id"]
|
||||
case_name = f"{index}_{document_id}" if document_id is not None else f"case_{index}"
|
||||
return Case(
|
||||
name=case_name,
|
||||
inputs=doc["question"],
|
||||
expected_output=doc["answer"],
|
||||
metadata={
|
||||
"document_id": str(document_id),
|
||||
"case_index": str(index),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
REPLIQA_SPEC = DatasetSpec(
|
||||
key="repliqa",
|
||||
db_filename="repliqa.lancedb",
|
||||
document_loader=load_repliqa_corpus,
|
||||
document_mapper=map_repliqa_document,
|
||||
qa_loader=load_repliqa_corpus,
|
||||
qa_case_builder=build_repliqa_case,
|
||||
retrieval_loader=load_repliqa_corpus,
|
||||
retrieval_mapper=map_repliqa_retrieval,
|
||||
retrieval_evaluator=MRREvaluator(),
|
||||
)
|
||||
|
|
@ -1,21 +1,15 @@
|
|||
from evaluations.evaluators.citation import (
|
||||
CitationMAPEvaluator,
|
||||
CitationMRREvaluator,
|
||||
)
|
||||
from evaluations.evaluators.citation import CitationMAPEvaluator
|
||||
from evaluations.evaluators.judge import (
|
||||
ANSWER_EQUIVALENCE_RUBRIC,
|
||||
LLMJudge,
|
||||
LLMJudgeResponseSchema,
|
||||
)
|
||||
from evaluations.evaluators.map import MAPEvaluator
|
||||
from evaluations.evaluators.mrr import MRREvaluator
|
||||
|
||||
__all__ = [
|
||||
"ANSWER_EQUIVALENCE_RUBRIC",
|
||||
"CitationMAPEvaluator",
|
||||
"CitationMRREvaluator",
|
||||
"LLMJudge",
|
||||
"LLMJudgeResponseSchema",
|
||||
"MAPEvaluator",
|
||||
"MRREvaluator",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -13,35 +13,13 @@ def _relevant_uris(ctx: EvaluatorContext) -> set[str]:
|
|||
return set(ctx.metadata.get("relevant_uris", []))
|
||||
|
||||
|
||||
@dataclass
|
||||
class CitationMRREvaluator(Evaluator):
|
||||
"""Reciprocal rank over the URIs the skill cited via the `cite` tool.
|
||||
|
||||
Reads ``cited_uris`` from ``ctx.attributes`` (recorded during the task run
|
||||
via :func:`pydantic_evals.set_eval_attribute`) and ``relevant_uris`` from
|
||||
``ctx.metadata``. Returns ``1.0/rank`` of the first cited URI that is in
|
||||
the relevant set, or ``0.0`` if none match.
|
||||
|
||||
Use for single-document datasets, mirroring :class:`MRREvaluator`.
|
||||
"""
|
||||
|
||||
def get_default_evaluation_name(self) -> str:
|
||||
return "cited_mrr"
|
||||
|
||||
def evaluate(self, ctx: EvaluatorContext) -> float:
|
||||
relevant = _relevant_uris(ctx)
|
||||
for rank, uri in enumerate(_cited_uris(ctx), start=1):
|
||||
if uri in relevant:
|
||||
return 1.0 / rank
|
||||
return 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class CitationMAPEvaluator(Evaluator):
|
||||
"""Average precision over the URIs the skill cited via the `cite` tool.
|
||||
|
||||
Same input shape as :class:`CitationMRREvaluator`; use for multi-document
|
||||
datasets, mirroring :class:`MAPEvaluator`.
|
||||
Reads ``cited_uris`` from ``ctx.attributes`` (recorded during the task run
|
||||
via :func:`pydantic_evals.set_eval_attribute`) and ``relevant_uris`` from
|
||||
``ctx.metadata``.
|
||||
"""
|
||||
|
||||
def get_default_evaluation_name(self) -> str:
|
||||
|
|
|
|||
|
|
@ -1,37 +0,0 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
|
||||
|
||||
|
||||
@dataclass
|
||||
class MRREvaluator(Evaluator):
|
||||
"""
|
||||
Mean Reciprocal Rank evaluator for single-document retrieval.
|
||||
|
||||
MRR = 1/rank where rank is the position of the first relevant document.
|
||||
Returns 0 if no relevant document is found.
|
||||
|
||||
Appropriate for retrieval tasks where each query has exactly one relevant document.
|
||||
"""
|
||||
|
||||
def evaluate(self, ctx: EvaluatorContext) -> float:
|
||||
"""
|
||||
Calculate reciprocal rank for a single query.
|
||||
|
||||
Expected context:
|
||||
- ctx.metadata['relevant_uris']: set/list of relevant document URIs
|
||||
- ctx.output: list of retrieved document URIs (ordered by rank)
|
||||
|
||||
Returns:
|
||||
float: 1/rank of first relevant doc, or 0.0 if not found
|
||||
"""
|
||||
if ctx.metadata is None:
|
||||
return 0.0
|
||||
relevant_uris = set(ctx.metadata.get("relevant_uris", []))
|
||||
retrieved_uris = ctx.output
|
||||
|
||||
for rank, uri in enumerate(retrieved_uris, start=1):
|
||||
if uri in relevant_uris:
|
||||
return 1.0 / rank
|
||||
|
||||
return 0.0
|
||||
|
|
@ -74,19 +74,19 @@ class TestBuildExperimentMetadata:
|
|||
|
||||
class TestResolveDataset:
|
||||
def test_valid_dataset(self) -> None:
|
||||
spec = _resolve_dataset("repliqa")
|
||||
assert spec.key == "repliqa"
|
||||
spec = _resolve_dataset("wix")
|
||||
assert spec.key == "wix"
|
||||
|
||||
def test_case_insensitive(self) -> None:
|
||||
spec = _resolve_dataset("REPLIQA")
|
||||
assert spec.key == "repliqa"
|
||||
spec = _resolve_dataset("WIX")
|
||||
assert spec.key == "wix"
|
||||
|
||||
def test_unknown_dataset_raises(self) -> None:
|
||||
with pytest.raises(typer.BadParameter, match="Unknown dataset 'nonexistent'"):
|
||||
_resolve_dataset("nonexistent")
|
||||
|
||||
def test_error_lists_valid_datasets(self) -> None:
|
||||
with pytest.raises(typer.BadParameter, match="repliqa"):
|
||||
with pytest.raises(typer.BadParameter, match="wix"):
|
||||
_resolve_dataset("nonexistent")
|
||||
|
||||
|
||||
|
|
@ -325,13 +325,6 @@ class TestRunQaBenchmarkSkillTarget:
|
|||
|
||||
|
||||
class TestCitationEvaluatorWiring:
|
||||
def test_returns_mrr_twin_for_mrr_evaluator(self) -> None:
|
||||
from evaluations.benchmark import _citation_evaluator_for
|
||||
from evaluations.evaluators import CitationMRREvaluator, MRREvaluator
|
||||
|
||||
result = _citation_evaluator_for(MRREvaluator())
|
||||
assert isinstance(result, CitationMRREvaluator)
|
||||
|
||||
def test_returns_map_twin_for_map_evaluator(self) -> None:
|
||||
from evaluations.benchmark import _citation_evaluator_for
|
||||
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
|
||||
|
|
@ -351,7 +344,7 @@ class TestAttachRelevantUris:
|
|||
|
||||
from evaluations.benchmark import _attach_relevant_uris
|
||||
from evaluations.config import RetrievalSample
|
||||
from evaluations.evaluators import MRREvaluator
|
||||
from evaluations.evaluators import MAPEvaluator
|
||||
|
||||
cases: list[Case[str, str, dict]] = [
|
||||
Case(name="c1", inputs="What is X?", expected_output="X is a thing"),
|
||||
|
|
@ -382,7 +375,7 @@ class TestAttachRelevantUris:
|
|||
retrieval_mapper=lambda d: RetrievalSample(
|
||||
question=d["q"], expected_uris=d["uris"]
|
||||
),
|
||||
retrieval_evaluator=MRREvaluator(),
|
||||
retrieval_evaluator=MAPEvaluator(),
|
||||
)
|
||||
|
||||
_attach_relevant_uris(cases, spec, limit=None)
|
||||
|
|
|
|||
|
|
@ -1,9 +1,6 @@
|
|||
from unittest.mock import MagicMock
|
||||
|
||||
from evaluations.evaluators.citation import (
|
||||
CitationMAPEvaluator,
|
||||
CitationMRREvaluator,
|
||||
)
|
||||
from evaluations.evaluators.citation import CitationMAPEvaluator
|
||||
|
||||
|
||||
def _ctx(cited: list[str], relevant: list[str]) -> MagicMock:
|
||||
|
|
@ -13,41 +10,6 @@ def _ctx(cited: list[str], relevant: list[str]) -> MagicMock:
|
|||
return ctx
|
||||
|
||||
|
||||
class TestCitationMRREvaluator:
|
||||
def setup_method(self) -> None:
|
||||
self.evaluator = CitationMRREvaluator()
|
||||
|
||||
def test_first_citation_is_relevant(self) -> None:
|
||||
assert self.evaluator.evaluate(_ctx(["a", "b"], ["a"])) == 1.0
|
||||
|
||||
def test_second_citation_is_relevant(self) -> None:
|
||||
assert self.evaluator.evaluate(_ctx(["a", "b"], ["b"])) == 0.5
|
||||
|
||||
def test_no_citations(self) -> None:
|
||||
assert self.evaluator.evaluate(_ctx([], ["a"])) == 0.0
|
||||
|
||||
def test_no_relevant(self) -> None:
|
||||
assert self.evaluator.evaluate(_ctx(["a"], [])) == 0.0
|
||||
|
||||
def test_no_matches(self) -> None:
|
||||
assert self.evaluator.evaluate(_ctx(["a", "b"], ["c"])) == 0.0
|
||||
|
||||
def test_metadata_none(self) -> None:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = None
|
||||
ctx.attributes = {"cited_uris": ["a"]}
|
||||
assert self.evaluator.evaluate(ctx) == 0.0
|
||||
|
||||
def test_attribute_missing(self) -> None:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = {"relevant_uris": ["a"]}
|
||||
ctx.attributes = {}
|
||||
assert self.evaluator.evaluate(ctx) == 0.0
|
||||
|
||||
def test_evaluation_name(self) -> None:
|
||||
assert self.evaluator.get_default_evaluation_name() == "cited_mrr"
|
||||
|
||||
|
||||
class TestCitationMAPEvaluator:
|
||||
def setup_method(self) -> None:
|
||||
self.evaluator = CitationMAPEvaluator()
|
||||
|
|
|
|||
|
|
@ -1,11 +1,5 @@
|
|||
from pathlib import Path
|
||||
|
||||
from evaluations.datasets.hotpotqa import (
|
||||
build_hotpotqa_case,
|
||||
extract_unique_documents,
|
||||
map_hotpotqa_document,
|
||||
map_hotpotqa_retrieval,
|
||||
)
|
||||
from evaluations.datasets.open_rag_bench import (
|
||||
build_orb_case,
|
||||
download_pdf,
|
||||
|
|
@ -13,11 +7,6 @@ from evaluations.datasets.open_rag_bench import (
|
|||
map_orb_document,
|
||||
map_orb_retrieval,
|
||||
)
|
||||
from evaluations.datasets.repliqa import (
|
||||
build_repliqa_case,
|
||||
map_repliqa_document,
|
||||
map_repliqa_retrieval,
|
||||
)
|
||||
from evaluations.datasets.wix import (
|
||||
build_wix_case,
|
||||
map_wix_document,
|
||||
|
|
@ -25,50 +14,6 @@ from evaluations.datasets.wix import (
|
|||
)
|
||||
|
||||
|
||||
class TestRepliqa:
|
||||
def test_map_document(self) -> None:
|
||||
doc = {"document_id": "doc-42", "document_extracted": "Some content here."}
|
||||
payload = map_repliqa_document(doc)
|
||||
assert payload.uri == "doc-42"
|
||||
assert payload.content == "Some content here."
|
||||
|
||||
def test_map_retrieval(self) -> None:
|
||||
doc = {
|
||||
"question": "What happened?",
|
||||
"answer": "Something happened.",
|
||||
"document_id": "doc-42",
|
||||
}
|
||||
sample = map_repliqa_retrieval(doc)
|
||||
assert sample is not None
|
||||
assert sample.question == "What happened?"
|
||||
assert sample.expected_uris == ("doc-42",)
|
||||
|
||||
def test_map_retrieval_skips_unanswerable(self) -> None:
|
||||
doc = {
|
||||
"question": "What?",
|
||||
"answer": "The answer is not found in the document.",
|
||||
"document_id": "doc-1",
|
||||
}
|
||||
assert map_repliqa_retrieval(doc) is None
|
||||
|
||||
def test_build_case(self) -> None:
|
||||
doc = {
|
||||
"document_id": "doc-7",
|
||||
"question": "Why?",
|
||||
"answer": "Because.",
|
||||
}
|
||||
case = build_repliqa_case(3, doc)
|
||||
assert case.name == "3_doc-7"
|
||||
assert case.inputs == "Why?"
|
||||
assert case.expected_output == "Because."
|
||||
assert case.metadata == {"document_id": "doc-7", "case_index": "3"}
|
||||
|
||||
def test_build_case_none_document_id(self) -> None:
|
||||
doc = {"document_id": None, "question": "Q?", "answer": "A."}
|
||||
case = build_repliqa_case(1, doc)
|
||||
assert case.name == "case_1"
|
||||
|
||||
|
||||
class TestWix:
|
||||
def test_map_document_with_all_fields(self) -> None:
|
||||
doc = {
|
||||
|
|
@ -136,78 +81,6 @@ class TestWix:
|
|||
assert case.name == "case_1"
|
||||
|
||||
|
||||
class TestHotpotQA:
|
||||
def test_map_document(self) -> None:
|
||||
doc = {"title": "Albert Einstein", "content": "Was a physicist."}
|
||||
payload = map_hotpotqa_document(doc)
|
||||
assert payload.uri == "Albert Einstein"
|
||||
assert payload.content == "Was a physicist."
|
||||
assert payload.title == "Albert Einstein"
|
||||
|
||||
def test_map_retrieval(self) -> None:
|
||||
doc = {
|
||||
"question": "Who was Einstein?",
|
||||
"supporting_facts": {"title": ["Albert Einstein", "Physics"]},
|
||||
}
|
||||
sample = map_hotpotqa_retrieval(doc)
|
||||
assert sample is not None
|
||||
assert sample.expected_uris == ("Albert Einstein", "Physics")
|
||||
|
||||
def test_map_retrieval_deduplicates_titles(self) -> None:
|
||||
doc = {
|
||||
"question": "Q?",
|
||||
"supporting_facts": {"title": ["A", "B", "A"]},
|
||||
}
|
||||
sample = map_hotpotqa_retrieval(doc)
|
||||
assert sample is not None
|
||||
assert sample.expected_uris == ("A", "B")
|
||||
|
||||
def test_map_retrieval_no_titles(self) -> None:
|
||||
doc = {"question": "Q?", "supporting_facts": {"title": []}}
|
||||
assert map_hotpotqa_retrieval(doc) is None
|
||||
|
||||
def test_build_case(self) -> None:
|
||||
doc = {
|
||||
"id": "abc123",
|
||||
"question": "What is X?",
|
||||
"answer": "X is Y.",
|
||||
"type": "comparison",
|
||||
"level": "hard",
|
||||
}
|
||||
case = build_hotpotqa_case(5, doc)
|
||||
assert case.name == "5_abc123"
|
||||
assert case.inputs == "What is X?"
|
||||
assert case.expected_output == "X is Y."
|
||||
assert case.metadata == {
|
||||
"question_id": "abc123",
|
||||
"type": "comparison",
|
||||
"level": "hard",
|
||||
"case_index": "5",
|
||||
}
|
||||
|
||||
def test_extract_unique_documents(self) -> None:
|
||||
# Simulate a minimal dataset with context
|
||||
dataset = [
|
||||
{
|
||||
"context": {
|
||||
"title": ["Doc A", "Doc B"],
|
||||
"sentences": [["Sentence 1."], ["Sentence 2.", " More."]],
|
||||
}
|
||||
},
|
||||
{
|
||||
"context": {
|
||||
"title": ["Doc A", "Doc C"],
|
||||
"sentences": [["Dupe."], ["Sentence 3."]],
|
||||
}
|
||||
},
|
||||
]
|
||||
docs = extract_unique_documents(dataset) # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
|
||||
assert len(docs) == 3
|
||||
titles = [d["title"] for d in docs]
|
||||
assert titles == ["Doc A", "Doc B", "Doc C"]
|
||||
assert docs[1]["content"] == "Sentence 2. More."
|
||||
|
||||
|
||||
class TestOpenRAGBench:
|
||||
def test_map_document(self, tmp_path: Path) -> None:
|
||||
# Pre-create a cached PDF
|
||||
|
|
|
|||
|
|
@ -3,54 +3,6 @@ from unittest.mock import MagicMock
|
|||
import pytest
|
||||
|
||||
from evaluations.evaluators.map import MAPEvaluator
|
||||
from evaluations.evaluators.mrr import MRREvaluator
|
||||
|
||||
|
||||
class TestMRREvaluator:
|
||||
def setup_method(self) -> None:
|
||||
self.evaluator = MRREvaluator()
|
||||
|
||||
def _make_ctx(
|
||||
self, relevant_uris: list[str], retrieved_uris: list[str]
|
||||
) -> MagicMock:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = {"relevant_uris": relevant_uris}
|
||||
ctx.output = retrieved_uris
|
||||
return ctx
|
||||
|
||||
def test_first_result_relevant(self) -> None:
|
||||
ctx = self._make_ctx(["doc1"], ["doc1", "doc2", "doc3"])
|
||||
assert self.evaluator.evaluate(ctx) == 1.0
|
||||
|
||||
def test_second_result_relevant(self) -> None:
|
||||
ctx = self._make_ctx(["doc2"], ["doc1", "doc2", "doc3"])
|
||||
assert self.evaluator.evaluate(ctx) == 0.5
|
||||
|
||||
def test_third_result_relevant(self) -> None:
|
||||
ctx = self._make_ctx(["doc3"], ["doc1", "doc2", "doc3"])
|
||||
assert self.evaluator.evaluate(ctx) == pytest.approx(1 / 3)
|
||||
|
||||
def test_no_relevant_found(self) -> None:
|
||||
ctx = self._make_ctx(["doc_x"], ["doc1", "doc2", "doc3"])
|
||||
assert self.evaluator.evaluate(ctx) == 0.0
|
||||
|
||||
def test_empty_retrieved(self) -> None:
|
||||
ctx = self._make_ctx(["doc1"], [])
|
||||
assert self.evaluator.evaluate(ctx) == 0.0
|
||||
|
||||
def test_multiple_relevant_returns_first_match(self) -> None:
|
||||
ctx = self._make_ctx(["doc2", "doc3"], ["doc1", "doc2", "doc3"])
|
||||
assert self.evaluator.evaluate(ctx) == 0.5
|
||||
|
||||
def test_none_metadata(self) -> None:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = None
|
||||
ctx.output = ["doc1"]
|
||||
assert self.evaluator.evaluate(ctx) == 0.0
|
||||
|
||||
def test_empty_relevant_uris(self) -> None:
|
||||
ctx = self._make_ctx([], ["doc1", "doc2"])
|
||||
assert self.evaluator.evaluate(ctx) == 0.0
|
||||
|
||||
|
||||
class TestMAPEvaluator:
|
||||
|
|
|
|||
|
|
@ -71,7 +71,6 @@ members = ["haiku_rag_slim", "evaluations"]
|
|||
dev = [
|
||||
"haiku.rag-evals",
|
||||
"haiku.rag-slim[ingester]",
|
||||
"datasets>=4.8.4",
|
||||
"zensical",
|
||||
"pre-commit>=4.5.1",
|
||||
"pydantic-ai-slim[anthropic]",
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
# Prevent tests from loading user's local haiku.rag.yaml by setting env var
|
||||
# to a test config file BEFORE any haiku.rag imports.
|
||||
|
|
@ -17,14 +18,12 @@ embeddings:
|
|||
vector_dim: 2560
|
||||
""")
|
||||
os.environ["HAIKU_RAG_CONFIG_PATH"] = str(_test_config_path)
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
|
||||
import pydantic_ai.models # noqa: E402
|
||||
import pytest # noqa: E402
|
||||
import yaml # noqa: E402
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from datasets import Dataset
|
||||
from vcr import VCR
|
||||
|
||||
setattr(pydantic_ai.models, "ALLOW_MODEL_REQUESTS", False)
|
||||
|
|
@ -32,18 +31,10 @@ logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def qa_corpus() -> "Dataset":
|
||||
from datasets import Dataset, load_dataset, load_from_disk
|
||||
|
||||
ds_path = Path(__file__).parent / "data" / "dataset"
|
||||
ds_path.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
return cast(Dataset, load_from_disk(ds_path))
|
||||
except FileNotFoundError:
|
||||
ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]
|
||||
corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
|
||||
corpus.save_to_disk(ds_path)
|
||||
return corpus
|
||||
def qa_corpus() -> list[dict[str, str]]:
|
||||
corpus_path = Path(__file__).parent / "data" / "qa_corpus.json"
|
||||
with open(corpus_path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
|
|||
Binary file not shown.
|
|
@ -1,106 +0,0 @@
|
|||
{
|
||||
"builder_name": "parquet",
|
||||
"citation": "",
|
||||
"config_name": "default",
|
||||
"dataset_name": "repliqa",
|
||||
"dataset_size": 648387648,
|
||||
"description": "",
|
||||
"download_checksums": {
|
||||
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_0-00000-of-00001.parquet": {
|
||||
"num_bytes": 20855947,
|
||||
"checksum": null
|
||||
},
|
||||
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_1-00000-of-00001.parquet": {
|
||||
"num_bytes": 20903799,
|
||||
"checksum": null
|
||||
},
|
||||
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_2-00000-of-00001.parquet": {
|
||||
"num_bytes": 20876034,
|
||||
"checksum": null
|
||||
},
|
||||
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_3-00000-of-00001.parquet": {
|
||||
"num_bytes": 20919214,
|
||||
"checksum": null
|
||||
},
|
||||
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_4-00000-of-00001.parquet": {
|
||||
"num_bytes": 20866527,
|
||||
"checksum": null
|
||||
}
|
||||
},
|
||||
"download_size": 104421521,
|
||||
"features": {
|
||||
"document_id": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
},
|
||||
"document_topic": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
},
|
||||
"document_path": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
},
|
||||
"document_extracted": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
},
|
||||
"question_id": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
},
|
||||
"question": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
},
|
||||
"answer": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
},
|
||||
"long_answer": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
}
|
||||
},
|
||||
"homepage": "",
|
||||
"license": "",
|
||||
"size_in_bytes": 752809169,
|
||||
"splits": {
|
||||
"repliqa_0": {
|
||||
"name": "repliqa_0",
|
||||
"num_bytes": 129579219,
|
||||
"num_examples": 17955,
|
||||
"dataset_name": "repliqa"
|
||||
},
|
||||
"repliqa_1": {
|
||||
"name": "repliqa_1",
|
||||
"num_bytes": 129844840,
|
||||
"num_examples": 17955,
|
||||
"dataset_name": "repliqa"
|
||||
},
|
||||
"repliqa_2": {
|
||||
"name": "repliqa_2",
|
||||
"num_bytes": 129338522,
|
||||
"num_examples": 17955,
|
||||
"dataset_name": "repliqa"
|
||||
},
|
||||
"repliqa_3": {
|
||||
"name": "repliqa_3",
|
||||
"num_bytes": 129981017,
|
||||
"num_examples": 17955,
|
||||
"dataset_name": "repliqa"
|
||||
},
|
||||
"repliqa_4": {
|
||||
"name": "repliqa_4",
|
||||
"num_bytes": 129644050,
|
||||
"num_examples": 17950,
|
||||
"dataset_name": "repliqa"
|
||||
}
|
||||
},
|
||||
"version": {
|
||||
"version_str": "0.0.0",
|
||||
"major": 0,
|
||||
"minor": 0,
|
||||
"patch": 0
|
||||
}
|
||||
}
|
||||
|
|
@ -1,13 +0,0 @@
|
|||
{
|
||||
"_data_files": [
|
||||
{
|
||||
"filename": "data-00000-of-00001.arrow"
|
||||
}
|
||||
],
|
||||
"_fingerprint": "2e29f63d782f12f0",
|
||||
"_format_columns": null,
|
||||
"_format_kwargs": {},
|
||||
"_format_type": null,
|
||||
"_output_all_columns": false,
|
||||
"_split": "repliqa_3"
|
||||
}
|
||||
252
tests/data/qa_corpus.json
Normal file
252
tests/data/qa_corpus.json
Normal file
File diff suppressed because one or more lines are too long
|
|
@ -1,5 +1,4 @@
|
|||
import pytest
|
||||
from datasets import Dataset
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
|
|
@ -7,7 +6,9 @@ from haiku.rag.store.models.chunk import Chunk, ChunkMetadata, SearchResult
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_chunk_repository_operations(
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Test ChunkRepository operations."""
|
||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
||||
# Get the first document from the corpus
|
||||
|
|
@ -46,7 +47,9 @@ async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_chunk_repository_pagination(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_chunk_repository_pagination(
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Test ChunkRepository pagination with get_by_document_id and count_by_document_id."""
|
||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
||||
# Get the first document from the corpus (should produce multiple chunks)
|
||||
|
|
@ -95,7 +98,7 @@ async def test_chunk_repository_pagination(qa_corpus: Dataset, temp_db_path):
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_chunking_pipeline(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_chunking_pipeline(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||
"""Test document chunking using client primitives."""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ from pathlib import Path
|
|||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from datasets import Dataset
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from haiku.rag.chunkers import get_chunker
|
||||
|
|
@ -18,7 +17,7 @@ def vcr_cassette_dir():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_chunker(qa_corpus: Dataset):
|
||||
async def test_local_chunker(qa_corpus: list[dict[str, str]]):
|
||||
"""Test DoclingLocalChunker with real document."""
|
||||
chunker = DoclingLocalChunker()
|
||||
doc_text = qa_corpus[0]["document_extracted"]
|
||||
|
|
@ -106,7 +105,7 @@ def test_tokenizer_cached_across_chunker_instances():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_chunker_hierarchical(qa_corpus: Dataset):
|
||||
async def test_local_chunker_hierarchical(qa_corpus: list[dict[str, str]]):
|
||||
"""Test DoclingLocalChunker with hierarchical chunking."""
|
||||
config = AppConfig()
|
||||
config.processing.chunker_type = "hierarchical"
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
from datasets import Dataset
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
|
|
@ -20,7 +19,7 @@ def vcr_cassette_dir():
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_client_document_crud(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_client_document_crud(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||
"""Test HaikuRAG CRUD operations for documents."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
# Get test data
|
||||
|
|
@ -123,7 +122,7 @@ async def test_client_resolve_document(temp_db_path):
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_client_update_document(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_client_update_document(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||
"""Test updating document with individual parameters."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
# Get test data
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import pytest
|
||||
from datasets import Dataset
|
||||
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
|
@ -8,7 +7,7 @@ from haiku.rag.store.repositories.document import DocumentRepository
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_list_excludes_content_by_default(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""list_all excludes content and docling_document by default."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
|
|
@ -34,7 +33,7 @@ async def test_document_list_excludes_content_by_default(
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_list_includes_content_when_requested(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""list_all returns content when include_content=True."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
|
|
@ -51,7 +50,7 @@ async def test_document_list_includes_content_when_requested(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||
"""Test listing documents with filter clause."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
doc_repo = DocumentRepository(store)
|
||||
|
|
@ -229,7 +228,7 @@ def test_get_page_images():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_docling_data_loads_only_docling_columns(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""get_docling_data returns docling blob without loading content."""
|
||||
import json
|
||||
|
|
@ -276,7 +275,9 @@ async def test_get_docling_data_loads_only_docling_columns(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_get_pages_data_loads_only_pages_column(
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""get_pages_data returns only page image data for a document."""
|
||||
import json
|
||||
|
||||
|
|
@ -309,7 +310,7 @@ async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_d
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_pages_data_none_for_markdown_document(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Markdown documents have no page images — get_pages_data returns None pages."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
|
|
@ -330,7 +331,7 @@ async def test_get_pages_data_none_for_markdown_document(
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_get_by_uri_with_special_characters(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Test get_by_uri handles URIs with special characters like single quotes."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ from pathlib import Path
|
|||
from typing import TypedDict
|
||||
|
||||
import pytest
|
||||
from datasets import Dataset
|
||||
|
||||
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||
|
||||
|
|
@ -19,10 +18,10 @@ class ChunkData(TypedDict):
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_full(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_rebuild_full(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||
"""Test full rebuild: converts, chunks, and embeds all documents."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
assert doc.id is not None
|
||||
assert doc.docling_document is not None
|
||||
|
||||
|
|
@ -49,10 +48,10 @@ async def test_rebuild_full(qa_corpus: Dataset, temp_db_path):
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_rebuild_embed_only(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||
"""Test embed-only rebuild: keeps chunks, only regenerates embeddings."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
assert doc.id is not None
|
||||
original_docling_json = doc.docling_document
|
||||
|
||||
|
|
@ -86,7 +85,7 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path):
|
|||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_embed_only_multi_doc_streams_via_staging(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Embed-only rebuild with multiple documents preserves chunks via staging.
|
||||
|
||||
|
|
@ -100,8 +99,8 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging(
|
|||
- the rebuild yields every document with chunks.
|
||||
"""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc1 = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc2 = await client.create_document(content=qa_corpus["document_extracted"][1])
|
||||
doc1 = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
doc2 = await client.create_document(content=qa_corpus[1]["document_extracted"])
|
||||
assert doc1.id is not None and doc2.id is not None
|
||||
|
||||
chunks_before_1 = await client.chunk_repository.get_by_document_id(doc1.id)
|
||||
|
|
@ -133,7 +132,9 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging(
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_rebuild_drops_leftover_staging_table(
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Staging table without marker is treated as partial phase 1 and dropped.
|
||||
|
||||
Simulates a phase-1 interruption by creating only the staging table (no
|
||||
|
|
@ -143,7 +144,7 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
|
|||
from haiku.rag.client.rebuild import _StagingChunkRecord
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
assert doc.id is not None
|
||||
|
||||
# Simulate a partial phase 1 (staging exists, marker absent).
|
||||
|
|
@ -167,7 +168,7 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
|
|||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_resumes_phase2_from_staging_after_crash(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Marker + staging present → phase 2 resumes from staging instead of
|
||||
redoing phase 1.
|
||||
|
|
@ -184,7 +185,7 @@ async def test_rebuild_resumes_phase2_from_staging_after_crash(
|
|||
)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
assert doc.id is not None
|
||||
original_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
assert original_chunks
|
||||
|
|
@ -280,7 +281,7 @@ async def test_rebuild_drops_orphan_marker(temp_db_path):
|
|||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Staging + marker from a prior embed-only crash → dropped on RECHUNK.
|
||||
|
||||
|
|
@ -295,7 +296,7 @@ async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
|
|||
)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
assert doc.id is not None
|
||||
|
||||
await client.store.db.create_table(
|
||||
|
|
@ -317,10 +318,12 @@ async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_rebuild_embed_only_skips_unchanged(
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Test embed-only rebuild skips chunks with unchanged embeddings."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
assert doc.id is not None
|
||||
|
||||
# Get embeddings before rebuild
|
||||
|
|
@ -354,7 +357,7 @@ async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_pa
|
|||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_embed_only_with_changed_vector_dim(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
qa_corpus: list[dict[str, str]], temp_db_path
|
||||
):
|
||||
"""Test embed-only rebuild when vector dimension changes.
|
||||
|
||||
|
|
@ -373,7 +376,7 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
|
|||
|
||||
# Step 1: Create a database with normal 2560-dim embeddings
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
assert doc.id is not None
|
||||
|
||||
chunks_before = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
|
|
@ -463,10 +466,10 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_rebuild_rechunk(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||
"""Test rechunk rebuild: re-chunks from content without accessing source files."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||
assert doc.id is not None
|
||||
assert doc.docling_document is not None
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import pytest
|
||||
from datasets import Dataset
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
|
|
@ -7,7 +6,7 @@ from haiku.rag.store.models import SearchResult
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
|
||||
async def test_search_qa_corpus(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||
"""Test that documents can be found by searching with their associated questions."""
|
||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
||||
# Load unique documents (limited to 10)
|
||||
|
|
|
|||
2
uv.lock
2
uv.lock
|
|
@ -1479,7 +1479,6 @@ tui = [
|
|||
|
||||
[package.dev-dependencies]
|
||||
dev = [
|
||||
{ name = "datasets" },
|
||||
{ name = "haiku-rag-evals" },
|
||||
{ name = "haiku-rag-slim", extra = ["ingester"] },
|
||||
{ name = "pre-commit" },
|
||||
|
|
@ -1506,7 +1505,6 @@ provides-extras = ["tui", "s3", "cross-encoder", "ingester"]
|
|||
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
{ name = "datasets", specifier = ">=4.8.4" },
|
||||
{ name = "haiku-rag-evals", editable = "evaluations" },
|
||||
{ name = "haiku-rag-slim", extras = ["ingester"], editable = "haiku_rag_slim" },
|
||||
{ name = "pre-commit", specifier = ">=4.5.1" },
|
||||
|
|
|
|||
Loading…
Reference in a new issue