diff --git a/README.md b/README.md index 37b320e..dc30448 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,7 @@ QueryGym implements the following query reformulation methods: | **ThinkQE** | Multi-round reasoning-based query expansion with corpus feedback | [Lei et al., 2025](https://arxiv.org/abs/2506.09260) | | **Query2E** | Query to entity/keyword expansion | [Jagerman et al., 2023](https://arxiv.org/abs/2305.03653)| | **ReFormeR** | Pattern-based, document-conditioned reformulation via learned transformation rules | [Bigdeli et al., 2026](https://arxiv.org/abs/2604.01417) | +| **HyDE** | Precise Zero-Shot Dense Retrieval without Relevance Labels | [Gao et al., 2022](https://arxiv.org/abs/2212.10496) | For detailed usage and parameters, see the [Methods Reference](https://querygym.readthedocs.io/en/latest/user-guide/methods-reference/). diff --git a/docs/about/changelog.md b/docs/about/changelog.md index 9b41995..335100d 100644 --- a/docs/about/changelog.md +++ b/docs/about/changelog.md @@ -5,6 +5,16 @@ All notable changes to querygym will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added +- HyDE reformulation method, from *Precise Zero-Shot Dense Retrieval without Relevance Labels* ([Gao et al., 2022](https://arxiv.org/abs/2212.10496)): generates 8 hypothetical documents per query and appends them to the query + - Task-specific prompts from the original implementation (`hyde.web_search.v1`, `hyde.scifact.v1`, `hyde.arguana.v1`, `hyde.trec_covid.v1`, `hyde.fiqa.v1`, `hyde.dbpedia_entity.v1`, `hyde.trec_news.v1`), selected with the `task` parameter + - `querygym.methods.hyde.task_for_dataset()` maps BEIR datasets to their prompt; every other dataset uses `web search`. The example pipeline uses it when `task` isn't set + - Paper sampling defaults in `defaults.yaml` (`temperature: 0.7`, `max_tokens: 512`, `top_p: 1.0`, `stop: ["\n\n\n"]`) + - `querygym run --output-format plain` writes one column per hypothetical document +- HyDE listed in README supported methods table + ## [0.3.0] - 2026-10-04 ### Added diff --git a/docs/api/methods.md b/docs/api/methods.md index 4dc8b0f..0afb6b6 100644 --- a/docs/api/methods.md +++ b/docs/api/methods.md @@ -135,3 +135,24 @@ reformulator = qg.create_reformulator( - `max_demo_len` (int): Optional word truncation for each passage before prompting (default: None) - `no_thinking` (bool): Prefill `` to disable reasoning traces (default: False) - `searcher`: Pre-configured searcher instance (recommended) + +### HyDE + +Hypothetical Document Embeddings, from *Precise Zero-Shot Dense Retrieval without Relevance Labels* ([Gao et al., 2022](https://arxiv.org/abs/2212.10496)), adapted to text: the LLM generates `n_docs` hypothetical documents with a task-specific prompt, and the reformulated query is the original query followed by the documents. + +```python +reformulator = qg.create_reformulator( + "hyde", + model="gpt-4.1", + params={"task": "scifact", "n_docs": 8}, + llm_config={"temperature": 0.7, "max_tokens": 512, "top_p": 1.0, "stop": ["\n\n\n"]} +) +``` + +**Requires Context:** No +**Parameters:** +- `task` (str): Prompt template: `web search`, `scifact`, `arguana`, `trec-covid`, `fiqa`, `dbpedia-entity`, `trec-news` (default: `web search`). An unknown task raises `ValueError` +- `n_docs` (int): Number of hypothetical documents to generate (default: 8) +- `repeat_query_weight` (int): Times the original query appears before the documents (default: 1) + +`querygym.methods.hyde.task_for_dataset(dataset_id)` returns the task for a dataset id (`web search` for any dataset without a dedicated prompt, or `None`). diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md index de904bf..d194d8b 100644 --- a/docs/getting-started/quickstart.md +++ b/docs/getting-started/quickstart.md @@ -49,6 +49,8 @@ Available methods: - `query2e` - Query to entity expansion - `csqe` - Context-based sentence extraction - `thinkqe` - Multi-round reasoning-based passage expansion +- `reformer` - Pattern-based, document-conditioned reformulation +- `hyde` - Hypothetical document generation ### 3. Reformulate Queries diff --git a/docs/user-guide/methods-reference.md b/docs/user-guide/methods-reference.md index 4609ee9..9591241 100644 --- a/docs/user-guide/methods-reference.md +++ b/docs/user-guide/methods-reference.md @@ -17,6 +17,7 @@ Complete reference guide for all query reformulation methods in QueryGym, includ - [CSQE](#csqe) - [ThinkQE](#thinkqe) - [ReFormeR](#reformer) + - [HyDE](#hyde) --- @@ -736,6 +737,85 @@ If the pattern selection response cannot be parsed as valid JSON, the method fal --- +### HyDE + +**Method Name:** `"hyde"` +**Requires Context:** No +**Description:** Hypothetical Document Embeddings, from *Precise Zero-Shot Dense Retrieval without Relevance Labels* ([Gao et al., 2022](https://arxiv.org/abs/2212.10496)). The LLM writes `n_docs` hypothetical documents that answer the query, using a task-specific prompt. The reformulated query is the original query followed by all generated documents: `query + doc_1 + ... + doc_8`. + +!!! note "Text adaptation" + The original HyDE averages the embeddings of the query and the hypothetical documents for dense retrieval. QueryGym emits text instead, so the output works with any retriever (BM25, SPLADE++, dense). The query appears once, matching its one-in-nine weight in the original embedding average. Dense encoders with a 512-token limit (e.g. BGE) only see the query and roughly the first document. + +#### Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `task` | str | `"web search"` | Prompt template to use (see [Tasks](#tasks)). An unknown task raises `ValueError` | +| `n_docs` | int | 8 | Number of hypothetical documents to generate (one LLM call each) | +| `repeat_query_weight` | int | 1 | Number of times the original query appears before the documents | + +#### LLM Configuration + +Defaults follow the HyDE paper: + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `temperature` | 0.7 | Sampling temperature | +| `max_tokens` | 512 | Maximum tokens per document | +| `top_p` | 1.0 | Nucleus sampling | +| `stop` | `["\n\n\n"]` | Stop sequences | + +!!! warning + `qg.create_reformulator` fills in `temperature=0.8` and `max_tokens=256` when `llm_config` doesn't set them. Pass `llm_config={"temperature": 0.7, "max_tokens": 512}` to use the paper settings. `querygym run` and the example pipeline config already use them. + +#### Tasks + +| `task` | Prompt ID | Prompt | +|--------|-----------|--------| +| `web search` | `hyde.web_search.v1` | Please write a passage to answer the question. | +| `scifact` | `hyde.scifact.v1` | Please write a scientific paper passage to support/refute the claim. | +| `arguana` | `hyde.arguana.v1` | Please write a counter argument for the passage. | +| `trec-covid` | `hyde.trec_covid.v1` | Please write a scientific paper passage to answer the question. | +| `fiqa` | `hyde.fiqa.v1` | Please write a financial article passage to answer the question. | +| `dbpedia-entity` | `hyde.dbpedia_entity.v1` | Please write a passage to answer the question. | +| `trec-news` | `hyde.trec_news.v1` | Please write a news passage about the topic. | + +The example pipeline (`examples/querygym_pyserini/`) picks the task from the dataset when `task` isn't set, using `querygym.methods.hyde.task_for_dataset()`: + +| Dataset | `task` | +|---------|--------| +| `beir-v1.0.0-scifact`, `-arguana`, `-trec-covid`, `-fiqa`, `-dbpedia-entity`, `-trec-news` | matching task | +| Everything else (MS MARCO, DL-HARD, BRIGHT, custom query files) | `web search` | + +An explicit `task` always wins over the dataset mapping. + +#### Usage Example + +```python +import querygym as qg + +reformulator = qg.create_reformulator( + "hyde", + model="gpt-4.1", + params={"task": "scifact", "n_docs": 8}, + llm_config={"temperature": 0.7, "max_tokens": 512, "top_p": 1.0, "stop": ["\n\n\n"]}, +) + +result = reformulator.reformulate(qg.QueryItem("q1", "0-dimensional biomaterials show inductive properties.")) +print(result.reformulated) # query + 8 hypothetical documents +print(result.metadata["hypothetical_docs"]) # the generated documents +``` + +#### How It Works + +1. **Prompt**: render the prompt for `task` with the query. +2. **Generate**: call the LLM `n_docs` times with the same prompt. Empty generations are dropped. +3. **Concatenate**: the original query (`repeat_query_weight` times) followed by the documents, with newlines flattened. If every generation is empty, the original query is returned. + +Metadata: `task`, `n_docs`, `hypothetical_docs`. + +--- + ## See Also - [API Reference](../api/methods.md) - Technical API documentation diff --git a/docs/user-guide/reformulation.md b/docs/user-guide/reformulation.md index 60dd61e..ce7c253 100644 --- a/docs/user-guide/reformulation.md +++ b/docs/user-guide/reformulation.md @@ -113,6 +113,19 @@ reformulator = qg.create_reformulator( ) ``` +### HyDE + +Generates 8 hypothetical documents with a task-specific prompt and appends them to the query. + +```python +reformulator = qg.create_reformulator( + "hyde", + model="gpt-4.1", + params={"task": "web search", "n_docs": 8}, + llm_config={"temperature": 0.7, "max_tokens": 512}, +) +``` + ## Method Comparison | Method | Requires Context | Type | Best For | @@ -126,6 +139,7 @@ reformulator = qg.create_reformulator( | query2e | No | Entity expansion | Entity queries | | csqe | Yes | Sentence extraction | Precision-focused | | thinkqe | Yes | Iterative reasoning | Multi-round feedback | +| hyde | No | Hypothetical documents | Zero-shot retrieval | ## Custom Parameters diff --git a/examples/querygym_pyserini/reformulate_queries.py b/examples/querygym_pyserini/reformulate_queries.py index a5e86c7..82fe05d 100755 --- a/examples/querygym_pyserini/reformulate_queries.py +++ b/examples/querygym_pyserini/reformulate_queries.py @@ -142,6 +142,12 @@ def reformulate_queries( retrieval_k = method_params.get('retrieval_k', 10) logging.info(f"Searcher configured with retrieval_k={retrieval_k}") + # HyDE: pick the task-specific prompt from the dataset unless one was given + if method == 'hyde' and 'task' not in method_params: + from querygym.methods.hyde import task_for_dataset + method_params['task'] = task_for_dataset(dataset_name) + logging.info(f"HyDE task (from dataset): {method_params['task']}") + # Create reformulator logging.info(f"Creating reformulator: {method}") logging.info(f"Model: {model}") diff --git a/examples/querygym_pyserini/reformulation_config.yaml b/examples/querygym_pyserini/reformulation_config.yaml index dfa6019..3f3355b 100644 --- a/examples/querygym_pyserini/reformulation_config.yaml +++ b/examples/querygym_pyserini/reformulation_config.yaml @@ -94,6 +94,20 @@ methods: max_tokens: 128 params: {} # No special parameters + # HyDE: Hypothetical document generation (query + 8 generated passages) + hyde: + enabled: true + model: "qwen2.5:7b" + llm: + temperature: 0.7 # HyDE paper sampling settings + max_tokens: 512 + top_p: 1.0 + stop: ["\n\n\n"] + params: + n_docs: 8 # Number of hypothetical documents to generate + # Note: task is picked from the dataset (e.g. scifact for beir-v1.0.0-scifact, + # web search otherwise). Set task here only to force one prompt for all datasets. + # LameR: Context-based passage synthesis (requires retrieval) lamer: enabled: true diff --git a/querygym/__init__.py b/querygym/__init__.py index 6f2ca52..f3bca0b 100644 --- a/querygym/__init__.py +++ b/querygym/__init__.py @@ -48,6 +48,7 @@ CSQE, ThinkQE, ReFormeR, + HyDE, ) # High-level runner @@ -241,6 +242,7 @@ def load_examples(path: str, **kwargs): "CSQE", "ThinkQE", "ReFormeR", + "HyDE", # High-level API "run_method", "build_llm", diff --git a/querygym/cli.py b/querygym/cli.py index e2701f6..57007ac 100644 --- a/querygym/cli.py +++ b/querygym/cli.py @@ -281,6 +281,15 @@ def write_plain_format(output_path): "reformulation_5", ] ) + elif method == "hyde": + # HyDE: qid \t passage_1 \t passage_2 \t ... \t passage_n + num_passages = ( + len(results[0].metadata.get("hypothetical_docs", [])) if results else 0 + ) + if not num_passages: + num_passages = 8 # Default n_docs + header = ["qid"] + [f"passage_{i+1}" for i in range(num_passages)] + w.writerow(header) elif method == "qa_expand": # QA-Expand: qid \t final_refined_query w.writerow(["qid", "refined_query"]) @@ -348,6 +357,10 @@ def write_plain_format(output_path): parts = reformulated.replace(original, "").strip().split() cleaned_parts = [clean_text(p) for p in parts] w.writerow([r.qid] + cleaned_parts) + elif method == "hyde": + # HyDE: qid \t passage_1 \t passage_2 \t ... \t passage_n + passages = r.metadata.get("hypothetical_docs", []) + w.writerow([r.qid] + [clean_text(p) for p in passages]) elif method == "query2doc": # Query2Doc: qid \t passage passage = r.metadata.get("pseudo_doc", "") diff --git a/querygym/config/defaults.yaml b/querygym/config/defaults.yaml index 968fece..984a5f2 100644 --- a/querygym/config/defaults.yaml +++ b/querygym/config/defaults.yaml @@ -71,5 +71,15 @@ reformer: context_docs: 3 # top-N docs fed to pattern selection (step 1) # patterns_path: null # override to load a custom patterns JSON file +hyde: + llm: + temperature: 0.7 # HyDE paper sampling settings + max_tokens: 512 + top_p: 1.0 + stop: ["\n\n\n"] + params: + task: "web search" # web search, scifact, arguana, trec-covid, fiqa, dbpedia-entity, trec-news + n_docs: 8 # Number of hypothetical documents to generate + seed: 42 retries: 2 \ No newline at end of file diff --git a/querygym/core/runner.py b/querygym/core/runner.py index 0ffa91a..3025117 100644 --- a/querygym/core/runner.py +++ b/querygym/core/runner.py @@ -16,6 +16,7 @@ from ..methods.csqe import CSQE from ..methods.thinkqe import ThinkQE from ..methods.reformer import ReFormeR +from ..methods.hyde import HyDE def build_llm(cfg: MethodConfig): diff --git a/querygym/methods/__init__.py b/querygym/methods/__init__.py index 7f2768f..2460988 100644 --- a/querygym/methods/__init__.py +++ b/querygym/methods/__init__.py @@ -8,3 +8,4 @@ from .csqe import CSQE from .thinkqe import ThinkQE from .reformer import ReFormeR +from .hyde import HyDE diff --git a/querygym/methods/hyde.py b/querygym/methods/hyde.py new file mode 100644 index 0000000..688ccf0 --- /dev/null +++ b/querygym/methods/hyde.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from typing import Dict, List, Optional + +from ..core.base import BaseReformulator, QueryItem, ReformulationResult +from ..core.registry import register_method + +DEFAULT_TASK = "web search" + +# Task-specific prompts from the original HyDE implementation (promptor.py). +# Mr. TyDi (multilingual) is not included. +TASK_PROMPTS: Dict[str, str] = { + "web search": "hyde.web_search.v1", + "scifact": "hyde.scifact.v1", + "arguana": "hyde.arguana.v1", + "trec-covid": "hyde.trec_covid.v1", + "fiqa": "hyde.fiqa.v1", + "dbpedia-entity": "hyde.dbpedia_entity.v1", + "trec-news": "hyde.trec_news.v1", +} + +# Datasets with a dedicated HyDE prompt. Every other dataset (MS MARCO, DL-HARD, +# BRIGHT, custom query files) uses the web search prompt. +DATASET_TASKS: Dict[str, str] = { + "beir-v1.0.0-scifact": "scifact", + "beir-v1.0.0-arguana": "arguana", + "beir-v1.0.0-trec-covid": "trec-covid", + "beir-v1.0.0-fiqa": "fiqa", + "beir-v1.0.0-dbpedia-entity": "dbpedia-entity", + "beir-v1.0.0-trec-news": "trec-news", +} + + +def task_for_dataset(dataset_id: Optional[str]) -> str: + """Return the HyDE task for a dataset id, defaulting to web search. + + Args: + dataset_id: Dataset id from dataset_registry.yaml, or None for custom query files + + Returns: + HyDE task name (a key of TASK_PROMPTS) + + Example: + >>> task_for_dataset("beir-v1.0.0-scifact") + 'scifact' + >>> task_for_dataset("bright-biology") + 'web search' + """ + return DATASET_TASKS.get(dataset_id or "", DEFAULT_TASK) + + +@register_method("hyde") +class HyDE(BaseReformulator): + """ + HyDE: Hypothetical Document Embeddings, from "Precise Zero-Shot Dense Retrieval + without Relevance Labels" (Gao et al., 2022). + + Text adaptation of HyDE: instead of averaging embeddings, the hypothetical + documents are concatenated with the original query, so the output works with + any retriever (BM25, SPLADE++, dense). + + Pipeline: + 1. Build the task-specific prompt (params["task"], default "web search") + 2. Call LLM n_docs times (default 8) to generate hypothetical documents + 3. Concatenate: query + doc_1 + doc_2 + ... + doc_n + """ + + VERSION = "1.0" + REQUIRES_CONTEXT = False + CONCATENATION_STRATEGY = "query_repeat_plus_generated" + # HyDE weights the query as one of n_docs + 1 inputs, so it appears once. + DEFAULT_QUERY_REPEATS = 1 + + def __init__(self, cfg, llm, pb): + super().__init__(cfg, llm, pb) + task = self.cfg.params.get("task", DEFAULT_TASK) + if task not in TASK_PROMPTS: + raise ValueError(f"Unknown HyDE task: {task!r}. Available: {list(TASK_PROMPTS.keys())}") + self.task = task + + def reformulate(self, q: QueryItem, contexts=None) -> ReformulationResult: + n_docs = int(self.cfg.params.get("n_docs", 8)) + temperature = float(self.cfg.llm.get("temperature", 0.7)) + max_tokens = int(self.cfg.llm.get("max_tokens", 512)) + top_p = float(self.cfg.llm.get("top_p", 1.0)) + stop = self.cfg.llm.get("stop", ["\n\n\n"]) + + msgs = self.prompts.render(TASK_PROMPTS[self.task], query=q.text) + + hypothetical_docs: List[str] = [] + for _ in range(n_docs): + out = self.llm.chat( + msgs, temperature=temperature, max_tokens=max_tokens, top_p=top_p, stop=stop + ) + doc = out.strip() + if doc: + hypothetical_docs.append(doc) + + if hypothetical_docs: + reformulated = self.concatenate_result(q.text, hypothetical_docs) + else: + reformulated = q.text + + return ReformulationResult( + q.qid, + q.text, + reformulated, + metadata={ + "task": self.task, + "n_docs": n_docs, + "hypothetical_docs": hypothetical_docs, + }, + ) diff --git a/querygym/prompt_bank.yaml b/querygym/prompt_bank.yaml index 35be643..8afaa81 100644 --- a/querygym/prompt_bank.yaml +++ b/querygym/prompt_bank.yaml @@ -550,3 +550,108 @@ Return only the reformulated query. notes: | ReFormeR step 2: applies the selected pattern's transformation rule to produce the final reformulated query. + +- id: hyde.web_search.v1 + method_family: hyde + version: 1 + introduced_by: "HyDE (Gao et al., 2022) — web search prompt" + license: "CC-BY-4.0" + authors: ["Luyu Gao", "Xueguang Ma", "Jimmy Lin", "Jamie Callan"] + tags: ["hypothetical-document", "generation", "expansion"] + template: + user: |- + Please write a passage to answer the question. + Question: {query} + Passage: + notes: | + HyDE hypothetical document prompt for the web search task, verbatim from the original promptor.py. No system message. + +- id: hyde.scifact.v1 + method_family: hyde + version: 1 + introduced_by: "HyDE (Gao et al., 2022) — SciFact prompt" + license: "CC-BY-4.0" + authors: ["Luyu Gao", "Xueguang Ma", "Jimmy Lin", "Jamie Callan"] + tags: ["hypothetical-document", "generation", "expansion"] + template: + user: |- + Please write a scientific paper passage to support/refute the claim. + Claim: {query} + Passage: + notes: | + HyDE hypothetical document prompt for the SciFact task, verbatim from the original promptor.py. No system message. + +- id: hyde.arguana.v1 + method_family: hyde + version: 1 + introduced_by: "HyDE (Gao et al., 2022) — ArguAna prompt" + license: "CC-BY-4.0" + authors: ["Luyu Gao", "Xueguang Ma", "Jimmy Lin", "Jamie Callan"] + tags: ["hypothetical-document", "generation", "expansion"] + template: + user: |- + Please write a counter argument for the passage. + Passage: {query} + Counter Argument: + notes: | + HyDE hypothetical document prompt for the ArguAna task, verbatim from the original promptor.py. No system message. + +- id: hyde.trec_covid.v1 + method_family: hyde + version: 1 + introduced_by: "HyDE (Gao et al., 2022) — TREC-COVID prompt" + license: "CC-BY-4.0" + authors: ["Luyu Gao", "Xueguang Ma", "Jimmy Lin", "Jamie Callan"] + tags: ["hypothetical-document", "generation", "expansion"] + template: + user: |- + Please write a scientific paper passage to answer the question. + Question: {query} + Passage: + notes: | + HyDE hypothetical document prompt for the TREC-COVID task, verbatim from the original promptor.py. No system message. + +- id: hyde.fiqa.v1 + method_family: hyde + version: 1 + introduced_by: "HyDE (Gao et al., 2022) — FiQA prompt" + license: "CC-BY-4.0" + authors: ["Luyu Gao", "Xueguang Ma", "Jimmy Lin", "Jamie Callan"] + tags: ["hypothetical-document", "generation", "expansion"] + template: + user: |- + Please write a financial article passage to answer the question. + Question: {query} + Passage: + notes: | + HyDE hypothetical document prompt for the FiQA task, verbatim from the original promptor.py. No system message. + +- id: hyde.dbpedia_entity.v1 + method_family: hyde + version: 1 + introduced_by: "HyDE (Gao et al., 2022) — DBPedia-Entity prompt" + license: "CC-BY-4.0" + authors: ["Luyu Gao", "Xueguang Ma", "Jimmy Lin", "Jamie Callan"] + tags: ["hypothetical-document", "generation", "expansion"] + template: + user: |- + Please write a passage to answer the question. + Question: {query} + Passage: + notes: | + HyDE hypothetical document prompt for the DBPedia-Entity task, verbatim from the original promptor.py. No system message. + +- id: hyde.trec_news.v1 + method_family: hyde + version: 1 + introduced_by: "HyDE (Gao et al., 2022) — TREC-News prompt" + license: "CC-BY-4.0" + authors: ["Luyu Gao", "Xueguang Ma", "Jimmy Lin", "Jamie Callan"] + tags: ["hypothetical-document", "generation", "expansion"] + template: + user: |- + Please write a news passage about the topic. + Topic: {query} + Passage: + notes: | + HyDE hypothetical document prompt for the TREC-News task, verbatim from the original promptor.py. No system message. diff --git a/tests/test_methods_hyde.py b/tests/test_methods_hyde.py new file mode 100644 index 0000000..dd1d126 --- /dev/null +++ b/tests/test_methods_hyde.py @@ -0,0 +1,223 @@ +from pathlib import Path + +import pytest + +import querygym as qg +from querygym.core.base import MethodConfig, QueryItem +from querygym.core.prompts import PromptBank +from querygym.core.registry import METHODS +from querygym.methods.hyde import DATASET_TASKS, TASK_PROMPTS, HyDE, task_for_dataset + +PROMPT_BANK = Path(__file__).parents[1] / "querygym" / "prompt_bank.yaml" + + +class RecordingLLM: + """Returns the queued outputs in order and records every call.""" + + def __init__(self, outputs=None): + self.outputs = list(outputs or []) + self.calls = [] + + def chat(self, messages, **kwargs): + self.calls.append((messages, kwargs)) + if self.outputs: + return self.outputs.pop(0) + return f"doc {len(self.calls)}" + + +def make_hyde(params=None, llm_cfg=None, llm=None): + cfg = MethodConfig(name="hyde", params=params or {}, llm={"model": "dummy", **(llm_cfg or {})}) + llm = llm or RecordingLLM() + return HyDE(cfg, llm, PromptBank(PROMPT_BANK)), llm + + +def test_generates_eight_docs_and_concatenates_query_once(): + meth, llm = make_hyde() + res = meth.reformulate(QueryItem("q1", "what causes diabetes")) + + assert len(llm.calls) == 8 + expected_docs = [f"doc {i}" for i in range(1, 9)] + assert res.reformulated == "what causes diabetes " + " ".join(expected_docs) + assert res.qid == "q1" + assert res.original == "what causes diabetes" + assert res.metadata == { + "task": "web search", + "n_docs": 8, + "hypothetical_docs": expected_docs, + } + + +def test_n_docs_param(): + meth, llm = make_hyde(params={"n_docs": 3}) + res = meth.reformulate(QueryItem("q1", "query")) + assert len(llm.calls) == 3 + assert res.reformulated == "query doc 1 doc 2 doc 3" + + +def test_repeat_query_weight(): + meth, _ = make_hyde(params={"n_docs": 1, "repeat_query_weight": 2}) + res = meth.reformulate(QueryItem("q1", "query")) + assert res.reformulated == "query query doc 1" + + +def test_default_prompt_is_web_search(): + meth, llm = make_hyde(params={"n_docs": 1}) + meth.reformulate(QueryItem("q1", "how long is a flea life cycle")) + messages = llm.calls[0][0] + assert messages == [ + { + "role": "user", + "content": "Please write a passage to answer the question.\n" + "Question: how long is a flea life cycle\nPassage:", + } + ] + + +@pytest.mark.parametrize( + "task, expected_start", + [ + ("web search", "Please write a passage to answer the question."), + ("scifact", "Please write a scientific paper passage to support/refute the claim."), + ("arguana", "Please write a counter argument for the passage."), + ("trec-covid", "Please write a scientific paper passage to answer the question."), + ("fiqa", "Please write a financial article passage to answer the question."), + ("dbpedia-entity", "Please write a passage to answer the question."), + ("trec-news", "Please write a news passage about the topic."), + ], +) +def test_task_selects_prompt(task, expected_start): + meth, llm = make_hyde(params={"task": task, "n_docs": 1}) + res = meth.reformulate(QueryItem("q1", "QUERY TEXT")) + content = llm.calls[0][0][0]["content"] + assert content.startswith(expected_start) + assert "QUERY TEXT" in content + assert res.metadata["task"] == task + + +def test_unknown_task_raises(): + with pytest.raises(ValueError, match="Unknown HyDE task: 'scifct'"): + make_hyde(params={"task": "scifct"}) + + +def test_sampling_params_passed_to_llm(): + meth, llm = make_hyde(params={"n_docs": 1}) + meth.reformulate(QueryItem("q1", "query")) + assert llm.calls[0][1] == { + "temperature": 0.7, + "max_tokens": 512, + "top_p": 1.0, + "stop": ["\n\n\n"], + } + + +def test_llm_config_overrides_sampling_params(): + meth, llm = make_hyde( + params={"n_docs": 1}, + llm_cfg={"temperature": 0.2, "max_tokens": 64, "top_p": 0.9, "stop": ["END"]}, + ) + meth.reformulate(QueryItem("q1", "query")) + assert llm.calls[0][1] == {"temperature": 0.2, "max_tokens": 64, "top_p": 0.9, "stop": ["END"]} + + +def test_empty_generations_are_dropped(): + llm = RecordingLLM(["first", " ", "", "second"]) + meth, _ = make_hyde(params={"n_docs": 4}, llm=llm) + res = meth.reformulate(QueryItem("q1", "query")) + assert res.reformulated == "query first second" + assert res.metadata["hypothetical_docs"] == ["first", "second"] + + +def test_all_empty_generations_fall_back_to_query(): + llm = RecordingLLM(["", " ", "\n"]) + meth, _ = make_hyde(params={"n_docs": 3}, llm=llm) + res = meth.reformulate(QueryItem("q1", "query")) + assert res.reformulated == "query" + assert res.metadata["hypothetical_docs"] == [] + + +def test_newlines_flattened(): + llm = RecordingLLM(["line one\nline two"]) + meth, _ = make_hyde(params={"n_docs": 1}, llm=llm) + res = meth.reformulate(QueryItem("q1", "query")) + assert res.reformulated == "query line one line two" + + +@pytest.mark.parametrize( + "dataset_id, task", + [ + ("beir-v1.0.0-scifact", "scifact"), + ("beir-v1.0.0-arguana", "arguana"), + ("beir-v1.0.0-trec-covid", "trec-covid"), + ("beir-v1.0.0-fiqa", "fiqa"), + ("beir-v1.0.0-dbpedia-entity", "dbpedia-entity"), + ("beir-v1.0.0-trec-news", "trec-news"), + ("msmarco-v1-passage.trecdl2019", "web search"), + ("msmarco-v1-passage.trecdl2020", "web search"), + ("msmarco-v1-passage.dlhard", "web search"), + ("bright-biology", "web search"), + ("bright-stackoverflow", "web search"), + ("my-custom-dataset", "web search"), + (None, "web search"), + ("", "web search"), + ], +) +def test_task_for_dataset(dataset_id, task): + assert task_for_dataset(dataset_id) == task + + +def test_mapped_tasks_have_prompts(): + pb = PromptBank(PROMPT_BANK) + assert set(DATASET_TASKS.values()) <= set(TASK_PROMPTS) + for prompt_id in TASK_PROMPTS.values(): + assert pb.get(prompt_id).method_family == "hyde" + + +def test_registered_and_created_via_factory(): + assert METHODS["hyde"] is HyDE + reformulator = qg.create_reformulator("hyde", model="dummy", llm_config={"api_key": "x"}) + assert isinstance(reformulator, HyDE) + + +def test_cli_bundled_defaults(): + from querygym import cli + + cfg = cli.resolve_config("hyde") + assert cfg["llm"]["temperature"] == 0.7 + assert cfg["llm"]["max_tokens"] == 512 + assert cfg["llm"]["top_p"] == 1.0 + assert cfg["llm"]["stop"] == ["\n\n\n"] + assert cfg["params"]["task"] == "web search" + assert cfg["params"]["n_docs"] == 8 + + +def test_cli_run_writes_concat_and_plain(tmp_path, monkeypatch): + from typer.testing import CliRunner + + from querygym import cli + from querygym.core import runner + + llm = RecordingLLM() + monkeypatch.setattr(runner, "build_llm", lambda cfg: llm) + queries = tmp_path / "queries.tsv" + queries.write_text("q1\tprime rate\n") + + result = CliRunner().invoke( + cli.app, + [ + "run", + "--method", + "hyde", + "--queries-tsv", + str(queries), + "--output-tsv", + str(tmp_path / "out.tsv"), + ], + ) + + assert result.exit_code == 0, result.output + docs = [f"doc {i}" for i in range(1, 9)] + concat = (tmp_path / "out_concat.tsv").read_text().splitlines() + assert concat == ["q1\tprime rate " + " ".join(docs)] + plain = (tmp_path / "out_plain.tsv").read_text().splitlines() + assert plain[0] == "\t".join(["qid"] + [f"passage_{i}" for i in range(1, 9)]) + assert plain[1] == "\t".join(["q1"] + docs) diff --git a/web/site/src/data/methods.json b/web/site/src/data/methods.json index 88c27c9..930ff0d 100644 --- a/web/site/src/data/methods.json +++ b/web/site/src/data/methods.json @@ -61,5 +61,12 @@ "tagline": "Query-to-entity / keyword expansion.", "paper": "Jagerman et al., 2023", "paper_url": "https://arxiv.org/abs/2305.03653" + }, + { + "id": "hyde", + "name": "HyDE", + "tagline": "Hypothetical documents written by the LLM.", + "paper": "Gao et al., 2022", + "paper_url": "https://arxiv.org/abs/2212.10496" } ]