From e82fde5309f6bcb5dfe737f15be99bc4ff5dbc51 Mon Sep 17 00:00:00 2001 From: Radin Hamidi Rad Date: Sun, 4 Oct 2026 17:42:05 -0400 Subject: [PATCH] query2e: remove chat-model framing from LLM output Chat-tuned models wrap the Q2E keyword list in introductory text, markdown headings, category labels and traces, which ended up in the expanded query. - Query2E cleans the LLM output before keyword parsing, controlled by a new clean_output param (default true). clean_output=false parses the raw output exactly as before. Prompts are unchanged. - Cleanup metadata (unclosed_think, truncated_at_query, cleanup_fallback, cleanup_empty) is recorded per query. - Add BaseReformulator.effective_params(); Query2E records clean_output: true when enabled, so cleaned runs get their own params_hash. - Pipeline records effective params, keeps empty recorded params, and reuses the recorded configuration when the reformulate step is skipped. - Add strip_think_trace() helper and --clean-output/--no-clean-output CLI flag. Query2E version 2.0. Fixes #42 Co-Authored-By: Claude Opus 5.5 --- docs/about/changelog.md | 11 + docs/api/core.md | 10 + docs/user-guide/cli.md | 2 +- docs/user-guide/methods-reference.md | 3 +- docs/user-guide/reformulation.md | 2 +- examples/querygym_pyserini/pipeline.py | 14 +- .../querygym_pyserini/reformulate_queries.py | 4 +- examples/querygym_pyserini/run_summary.py | 5 +- querygym/cli.py | 7 + querygym/core/base.py | 21 ++ querygym/core/utils.py | 26 +- querygym/methods/query2e.py | 328 +++++++++++++++++- reproducibility/tests/test_repro_schema.py | 60 ++++ tests/test_methods_query2e.py | 211 +++++++++++ 14 files changed, 688 insertions(+), 16 deletions(-) create mode 100644 tests/test_methods_query2e.py diff --git a/docs/about/changelog.md b/docs/about/changelog.md index 9d49a5a..753c1dc 100644 --- a/docs/about/changelog.md +++ b/docs/about/changelog.md @@ -5,6 +5,17 @@ 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 +- `BaseReformulator.effective_params()`: method parameters that identify a run's configuration +- `querygym.core.utils.strip_think_trace()` helper for removing `` reasoning traces +- `--clean-output/--no-clean-output` CLI flag + +### Changed +- Query2E removes chat-model framing (`` traces, introductory/closing sentences, markdown headings, category labels, term descriptions) from the LLM output before keyword parsing ([#42](https://github.com/ls3-lab/QueryGym/issues/42)). Controlled by the new `clean_output` parameter (default `True`); `clean_output=False` parses the raw output as before. Query2E method version 2.0 +- Query2E runs record `clean_output: true` in their method parameters when output cleanup is enabled + ## [0.2.0] - 2026-04-12 ### Added diff --git a/docs/api/core.md b/docs/api/core.md index 76df64d..2b743cf 100644 --- a/docs/api/core.md +++ b/docs/api/core.md @@ -108,6 +108,16 @@ for r in results: print(f"{r.qid}: {r.reformulated}") ``` +#### effective_params + +```python +params = reformulator.effective_params() +``` + +Method parameters that identify a run's configuration: the JSON-serializable subset of the configured params (runtime objects such as a searcher are dropped). Methods override it to record a parameter whose default changed, e.g. Query2E adds `clean_output: true` when output cleanup is enabled. + +**Returns:** `Dict[str, Any]` + #### concatenate_result ```python diff --git a/docs/user-guide/cli.md b/docs/user-guide/cli.md index e768138..fbc5f0e 100644 --- a/docs/user-guide/cli.md +++ b/docs/user-guide/cli.md @@ -1,6 +1,6 @@ # CLI -- `run` — produce reformulated TSV. +- `run` — produce reformulated TSV. For Query2E, `--no-clean-output` parses the raw LLM output instead of removing chat-model framing first. - `data-to-tsv` — convert any backend to `qidquery`. - `prompts-list` / `prompts-show` — inspect prompt bank. - `script-gen` — create a Pyserini + trec_eval bash script. diff --git a/docs/user-guide/methods-reference.md b/docs/user-guide/methods-reference.md index f2458aa..4609ee9 100644 --- a/docs/user-guide/methods-reference.md +++ b/docs/user-guide/methods-reference.md @@ -445,6 +445,7 @@ result = reformulator.reformulate(qg.QueryItem("q1", "machine learning")) | `mode` | str | `"zs"` | Mode: `"zs"`/`"zeroshot"` (zero-shot) or `"fs"`/`"fewshot"` (few-shot) | | `num_examples` | int | `4` | Number of few-shot examples (only for `mode="fs"`) | | `max_keywords` | int | `20` | Maximum number of keywords to extract | +| `clean_output` | bool | `True` | Remove chat-model framing from the LLM output before keyword parsing: `` traces, introductory/closing sentences, markdown headings, category labels and term descriptions. `False` parses the raw output | | `dataset_type` | str | `None` | Dataset type for few-shot: `"msmarco"`, `"beir"`, or `"generic"` | | `collection_path` | str | `None` | Path to collection file (for MS MARCO/generic) | | `train_queries_path` | str | `None` | Path to training queries file | @@ -491,7 +492,7 @@ result = reformulator.reformulate(qg.QueryItem("q1", "deep learning")) #### Output Format - **Concatenation:** `(query × 5) + keywords` -- **Metadata:** Includes `mode`, `keywords` list, `prompt_id`, `num_examples` (for few-shot) +- **Metadata:** Includes `mode`, `keywords` list, `prompt_id`, `num_examples` (for few-shot). With `clean_output`, also `clean_output` and, when applicable, `unclosed_think` (reasoning trace not closed), `truncated_at_query` (model continued with another few-shot example), `cleanup_fallback` (raw keyword list used) or `cleanup_empty` (no keywords found) --- diff --git a/docs/user-guide/reformulation.md b/docs/user-guide/reformulation.md index bcfaf31..60dd61e 100644 --- a/docs/user-guide/reformulation.md +++ b/docs/user-guide/reformulation.md @@ -75,7 +75,7 @@ results = reformulator.reformulate_batch(queries, contexts=contexts) ### Query2E -Query to entity expansion. +Query to keyword expansion. Chat-model framing (introductory sentences, markdown, `` traces) is removed from the LLM output before keyword parsing; pass `clean_output=False` to parse the raw output. ```python reformulator = qg.create_reformulator("query2e", model="gpt-4") diff --git a/examples/querygym_pyserini/pipeline.py b/examples/querygym_pyserini/pipeline.py index 4941390..e90078b 100755 --- a/examples/querygym_pyserini/pipeline.py +++ b/examples/querygym_pyserini/pipeline.py @@ -123,7 +123,19 @@ def run_pipeline( except Exception as e: logging.error(f"✗ Reformulation failed: {e}") raise - + elif (output_dir / 'reformulation_metadata.json').exists(): + # Reuse the configuration that produced the existing reformulated queries + with open(output_dir / 'reformulation_metadata.json') as f: + metadata = json.load(f) + recorded = metadata.get('reformulation', {}) + if (recorded.get('method'), recorded.get('model')) != (method, model): + raise ValueError( + f"Reformulated queries in {output_dir} were produced with " + f"method={recorded.get('method')}, model={recorded.get('model')}; " + f"got method={method}, model={model}" + ) + results['reformulation'] = metadata + # Step 2: Retrieve documents if 'retrieve' in steps: logging.info("\n" + "="*60) diff --git a/examples/querygym_pyserini/reformulate_queries.py b/examples/querygym_pyserini/reformulate_queries.py index f74d7eb..a5e86c7 100755 --- a/examples/querygym_pyserini/reformulate_queries.py +++ b/examples/querygym_pyserini/reformulate_queries.py @@ -181,8 +181,8 @@ def reformulate_queries( qg.DataLoader.save_queries(queries, output_queries_original, format='tsv') logging.info(f"Saved original queries: {output_queries_original}") - # Prepare method params for metadata (exclude non-serializable searcher object) - metadata_params = {k: v for k, v in method_params.items() if k != 'searcher'} + # Method params that identify the run (excludes runtime objects such as the searcher) + metadata_params = reformulator.effective_params() # Build dataset metadata - handle both registry-based and file-based inputs if dataset_name: diff --git a/examples/querygym_pyserini/run_summary.py b/examples/querygym_pyserini/run_summary.py index c34d9da..b771b27 100644 --- a/examples/querygym_pyserini/run_summary.py +++ b/examples/querygym_pyserini/run_summary.py @@ -76,7 +76,10 @@ def _build_v1_summary( 'implementation': f'pyserini:{searcher_class}', } - eff_method_params = reform_inner.get('method_params') or dict(method_params or {}) + if 'method_params' in reform_inner: + eff_method_params = dict(reform_inner['method_params']) + else: + eff_method_params = dict(method_params or {}) eff_method_params = { k: v for k, v in eff_method_params.items() if not callable(v) and not hasattr(v, '__dict__') or isinstance(v, (dict, list, str, int, float, bool, type(None))) diff --git a/querygym/cli.py b/querygym/cli.py index 2cc98b8..3cb0102 100644 --- a/querygym/cli.py +++ b/querygym/cli.py @@ -60,6 +60,11 @@ def run( mode: Optional[str] = typer.Option( None, "--mode", help="Mode for methods: 'zs' (zero-shot) or 'fs' (few-shot)" ), + clean_output: Optional[bool] = typer.Option( + None, + "--clean-output/--no-clean-output", + help="Query2E: remove chat-model framing from LLM output before keyword parsing (default: on)", + ), # Few-shot data loading options dataset_type: Optional[str] = typer.Option( None, "--dataset-type", help="Dataset type: 'msmarco', 'beir', or 'generic'" @@ -114,6 +119,8 @@ def replace_env_var(match): params["parallel"] = parallel if mode is not None: params["mode"] = mode + if clean_output is not None: + params["clean_output"] = clean_output if dataset_type is not None: params["dataset_type"] = dataset_type if collection_path is not None: diff --git a/querygym/core/base.py b/querygym/core/base.py index f65f245..4acbecf 100644 --- a/querygym/core/base.py +++ b/querygym/core/base.py @@ -1,4 +1,5 @@ from __future__ import annotations +import json from dataclasses import dataclass, field from typing import Dict, List, Optional, Any from tqdm import tqdm @@ -107,6 +108,26 @@ def __init__(self, cfg: MethodConfig, llm_client, prompt_resolver): self.llm = llm_client self.prompts = prompt_resolver + def effective_params(self) -> Dict[str, Any]: + """Method parameters that identify a run's configuration. + + Returns the JSON-serializable subset of the configured params (runtime + objects such as a searcher are dropped). Methods override this to record + a parameter whose default changed, so runs with different behavior are + distinguishable from runs recorded before the change. + + Returns: + Dict of parameter names to values + """ + params = {} + for k, v in self.cfg.params.items(): + try: + json.dumps(v) + except (TypeError, ValueError): + continue + params[k] = v + return params + def reformulate( self, q: QueryItem, contexts: Optional[List[str]] = None ) -> ReformulationResult: diff --git a/querygym/core/utils.py b/querygym/core/utils.py index c339c1f..f08d480 100644 --- a/querygym/core/utils.py +++ b/querygym/core/utils.py @@ -1,6 +1,30 @@ from __future__ import annotations import random -from typing import Optional +import re +from typing import Optional, Tuple + +_THINK_BLOCK = re.compile(r".*?", re.DOTALL) + + +def strip_think_trace(text: str) -> Tuple[str, bool]: + """Remove ... reasoning traces from model output. + + Args: + text: Raw model output + + Returns: + (text without reasoning traces, whether an unclosed was found). + When a block is never closed (generation ran out of tokens while + reasoning), everything from the opening tag on is dropped. When only the + closing tag is present (chat templates that put in the prompt), + everything up to it is dropped. + """ + text = _THINK_BLOCK.sub("", text) + if "" in text: + text = text.split("", 1)[1] + if "" in text: + return text.split("", 1)[0], True + return text, False def seed_everything(seed: Optional[int] = 42): diff --git a/querygym/methods/query2e.py b/querygym/methods/query2e.py index c2019ee..59e2bfa 100644 --- a/querygym/methods/query2e.py +++ b/querygym/methods/query2e.py @@ -1,10 +1,99 @@ from __future__ import annotations from ..core.base import BaseReformulator, QueryItem, ReformulationResult from ..core.registry import register_method -from typing import List, Tuple +from ..core.utils import strip_think_trace +from typing import Any, Dict, List, Tuple import random import os import csv +import json +import re + +# Patterns for removing the conversational and markdown framing that chat-tuned +# models wrap around a keyword list (see Query2E._strip_framing) +_HEADING = re.compile(r"^#{1,6}\s") +_RULE = re.compile(r"^[-*_=]{3,}$") +_FENCE = re.compile(r"^```|^~~~") +_TABLE_ROW = re.compile(r"^\|.*\|$") +_BULLET = re.compile(r"^([\-\*•▪→►✅✔☑]|\(?\d+[\).:]|[a-z]\))\s+") +_BOLD_LABEL = re.compile(r"^(\*\*|__)([^*_]+?)(\*\*|__)\s*(:?)\s*(.*)$") +_HTML_INLINE = re.compile(r"", re.IGNORECASE) +_HTML_BLOCK = re.compile(r"", re.IGNORECASE) +_QUERY_LINE = re.compile(r"^\**\s*query\s*\**\s*[:：]", re.IGNORECASE) +_KEYWORDS_LABEL = re.compile( + r"^\**\s*(keywords?|mots[- ]cl[ée]s|palabras clave|schl[üu]sselw[öo]rter|关键词|キーワード|" + r"ключевые слова)\s*(\([^)]*\))?\s*\**\s*[:：]\s*\**\s*", + re.IGNORECASE, +) +_CLOSER_PHRASES = ( + r"note\b|would you like|let me know|i hope|hope this|feel free|" + r"these (?:keywords|terms|should)|this list" +) +_CLOSER = re.compile(rf"^(?:{_CLOSER_PHRASES})", re.IGNORECASE) +_OPENER = re.compile( + r"^(?:sure|certainly|of course|okay|absolutely|here(?:'s|’s| is| are)|below (?:is|are|you)|" + rf"the following|i can help|to find|you (?:can|could|might|may)|if you|{_CLOSER_PHRASES})", + re.IGNORECASE, +) +_OPENER_OTHER_LANG = re.compile( + r"^(aqu[ií]|voici|ecco|hier (ist|sind)|a continuaci[oó]n|以下|下面|다음|voil[aà])", + re.IGNORECASE, +) +_LEXICON = re.compile(r"\b(keyword|keywords|list|query|phrase|phrases|terms?)\b", re.IGNORECASE) +_LIST_VERBS = r"are|include|includes|including|would be|be|use|using" +_VERB = re.compile(rf"\b(?:{_LIST_VERBS}|is|can|should)\b", re.IGNORECASE) +_LIST_TAIL = re.compile(rf"(?:[:：]|\b(?:{_LIST_VERBS}|such as)\b)\s*", re.IGNORECASE) +_PROSE_WORDS = set( + "you your i we these those can will would should could may might hope help helps let know " + "feel free need want like please if here below following above some relevant related based " + "broken".split() +) +_SENTENCE_END = (".", "。", "!", "！", "?", "？") + + +def _strip_emphasis(s: str) -> str: + """Remove markdown/HTML emphasis while keeping identifiers like __init__.""" + s = _HTML_BLOCK.sub("", _HTML_INLINE.sub("", s)) + s = re.sub(r"\*\*(.+?)\*\*", r"\1", s) + s = re.sub(r"(? bool: + """At least three short comma-separated items.""" + parts = [p.strip() for p in re.split(r"[,;，；、]", s) if p.strip()] + return len(parts) >= 3 and sum(len(p.split()) for p in parts) / len(parts) <= 6 + + +def _is_prose(s: str) -> bool: + """Heuristic: the line reads like a sentence rather than keywords.""" + tokens = re.findall(r"[a-zA-Z']+", s.lower()) + if sum(1 for t in tokens if t in _PROSE_WORDS) >= 2: + return True + if _LEXICON.search(" ".join(s.split()[:6])) and _VERB.search(s): + return True + if len(tokens) >= 12 and (s.endswith(_SENTENCE_END) or "," not in s): + return True + return len(tokens) >= 6 and s.endswith(_SENTENCE_END[2:]) and "," not in s + + +def _list_tail(s: str) -> str | None: + """Keyword list embedded at the end of a sentence ("... include a, b, c").""" + tail = None + for m in _LIST_TAIL.finditer(s): + candidate = s[m.end() :].strip().rstrip("".join(_SENTENCE_END)) + if _is_list_like(candidate) and not _is_prose(candidate): + tail = candidate + if tail is None: + return None + tail = re.sub(r"^(and|or)\s+", "", tail.strip()) + return re.sub(r",\s*(and|or)\s+", ", ", tail) + + +def _indent(line: str) -> int: + return len(line) - len(line.lstrip()) @register_method("query2e") @@ -18,6 +107,12 @@ class Query2E(BaseReformulator): Formula: (query × 5) + keywords + Output Cleanup (params["clean_output"], default True): + Chat-tuned models tend to wrap the keyword list in introductory and closing + sentences, markdown headings, category labels and reasoning traces. + With clean_output enabled this framing is removed before keyword parsing. + Set clean_output=False to parse the raw model output (behavior before 2.0). + Few-Shot Examples (3 ways to provide): 1. Via --ctx-jsonl CLI flag (JSONL file with {"query": "...", "passage": "..."} per line) 2. Via params["examples"] (list of {"query": "...", "passage": "..."} dicts) @@ -44,7 +139,7 @@ class Query2E(BaseReformulator): Note: MS MARCO env vars (MSMARCO_COLLECTION, etc.) are supported for backward compatibility. """ - VERSION = "1.0" + VERSION = "2.0" CONCATENATION_STRATEGY = "query_repeat_plus_generated" DEFAULT_QUERY_REPEATS = 5 @@ -226,7 +321,211 @@ def _format_examples(self, examples: List[Tuple[str, str]]) -> str: return examples_text - def _parse_keywords(self, raw_output: str) -> List[str]: + def _clean_output_enabled(self) -> bool: + return bool(self.cfg.params.get("clean_output", True)) + + def effective_params(self) -> Dict[str, Any]: + """Record clean_output when enabled; runs without it parsed the raw output.""" + params = super().effective_params() + if self._clean_output_enabled(): + params["clean_output"] = True + else: + params.pop("clean_output", None) + return params + + def _strip_framing(self, raw_output: str) -> Tuple[str, Dict[str, Any]]: + """ + Remove the conversational and markdown framing that chat-tuned models add + around a keyword list: + - reasoning traces + - introductory/closing sentences ("Here is a list of keywords...", "I hope this helps!") + - markdown headings, horizontal rules, code fences, category labels and emphasis + - "Term: description" explanations (the term is kept) + - further "Query:/Keywords:" pairs a model adds after its answer in few-shot mode + + Keyword lists, table cells and JSON string arrays are kept. + + Returns: + (cleaned text in a format accepted by _parse_keywords, cleanup metadata) + """ + meta: Dict[str, Any] = {} + if not raw_output: + return "", meta + + text, unclosed = strip_think_trace(raw_output) + if unclosed: + meta["unclosed_think"] = True + + lines = text.split("\n") + kept: List[str] = [] + seen_content = False + table_rows = 0 + keyword_columns: List[int] = [] + + for i, line in enumerate(lines): + s = line.strip() + if not s: + continue + if _HEADING.match(s) or _RULE.match(s) or _FENCE.match(s): + continue + + # Table: skip header and separator rows, keep the keyword cells. The + # keyword column is taken from the header, else every column but the first + if _TABLE_ROW.match(s): + cells = [c.strip() for c in s.strip("|").split("|")] + if all(re.fullmatch(r":?-{2,}:?", c) for c in cells): + continue + table_rows += 1 + if table_rows == 1: + keyword_columns = [ + j for j, c in enumerate(cells) if re.search(r"keyword|term|phrase", c, re.I) + ] + continue + if keyword_columns: + values = [cells[j] for j in keyword_columns if j < len(cells)] + else: + values = cells[1:] if len(cells) > 1 else cells + kept.append(", ".join(_strip_emphasis(c) for c in values if c)) + seen_content = True + continue + + # JSON array of strings + if s.startswith("[") and s.endswith("]"): + try: + items = json.loads(s) + except ValueError: + items = None + if isinstance(items, list) and all(isinstance(x, str) for x in items): + kept.append(", ".join(items)) + seen_content = True + continue + + # Few-shot continuation: the model started another example + if _QUERY_LINE.match(s): + if seen_content: + meta["truncated_at_query"] = True + break + continue + + # HTML lists: "
  • term
    • " is a list item + if _HTML_BLOCK.search(s): + html_item = re.match(r" _indent(line) + + # "Keywords: a, b, c" + label = _KEYWORDS_LABEL.match(body) + if label: + body = body[label.end() :].strip() + if body: + kept.append(_strip_emphasis(body)) + seen_content = True + continue + + # List item acting as a category header for a nested list + if ( + is_list_item + and next_is_nested + and _BULLET.match(next_line.strip()) + and not re.search(r"[,:：]", _strip_emphasis(body)) + ): + continue + + # Bold labels: "**Label:** a, b", "**Category**", "- **Term**: description" + bold = _BOLD_LABEL.match(_HTML_INLINE.sub("**", body)) + if bold: + name, colon_outside, rest = bold.group(2).strip(), bold.group(4), bold.group(5) + if name.endswith(":"): + if rest: + kept.append(_strip_emphasis(rest)) + seen_content = True + continue + if not rest and not is_list_item: + continue + if colon_outside or (rest and re.match(r"^[\-–—(]", rest)): + rest = rest.lstrip(":-–— ").strip() + if is_list_item and not _is_list_like(rest): + kept.append(name) + else: + kept.append(_strip_emphasis(rest) if rest else name) + seen_content = True + continue + + # Standalone "*Category*" line + if not is_list_item and re.fullmatch(r"\*[^*]+\*|_[^_]+_", body): + continue + + plain = _strip_emphasis(body) + + # "- term - explanation of the term" + if is_list_item and re.search(r"\s[-–—]\s", plain): + head, explanation = re.split(r"\s[-–—]\s", plain, maxsplit=1) + if len(explanation.split()) >= 4 or _is_prose(explanation): + plain = head.strip() + + # Lines ending with a colon are lead-ins or category labels + if plain.endswith((":", "：")): + if not is_list_item or next_is_nested: + continue + plain = plain.rstrip(":： ") + + if not is_list_item: + # A bare keyword list is kept even if it starts like a sentence ("okay google, ...") + if _is_list_like(plain) and not _VERB.search(plain) and not _CLOSER.match(plain): + kept.append(plain) + seen_content = True + continue + if ( + _OPENER.match(plain) + or _OPENER_OTHER_LANG.match(plain) + or _CLOSER.match(plain) + or _is_prose(plain) + ): + tail = None + if (_is_list_like(plain) or "," in plain) and not _CLOSER.match(plain): + tail = _list_tail(plain) + if tail: + kept.append(tail) + seen_content = True + continue + + kept.append(plain) + seen_content = True + + return "\n".join(kept), meta + + def _extract_keywords(self, raw_output: str) -> Tuple[List[str], Dict[str, Any]]: + """Clean the model output and parse it into keywords.""" + if not raw_output or not raw_output.strip(): + return [], {} + + cleaned, meta = self._strip_framing(raw_output) + terms = self._parse_keywords(cleaned, strict_numbering=True) + if terms or meta.get("unclosed_think"): + return terms, meta + + # Nothing survived cleanup: fall back to the output (without reasoning + # traces) only if it is itself a plain keyword list + raw_terms = self._parse_keywords(strip_think_trace(raw_output)[0], strict_numbering=True) + if ( + len(raw_terms) >= 2 + and all(len(t.split()) <= 5 for t in raw_terms) + and not any(_is_prose(t) for t in raw_terms) + ): + meta["cleanup_fallback"] = True + return raw_terms, meta + meta["cleanup_empty"] = True + return [], meta + + def _parse_keywords(self, raw_output: str, strict_numbering: bool = False) -> List[str]: """ Parse keywords from LLM output, handling various formats: - Comma-separated: "keyword1, keyword2, keyword3" @@ -234,11 +533,15 @@ def _parse_keywords(self, raw_output: str) -> List[str]: - Numbered lists: "1. keyword1\n2. keyword2" - Mixed formats + Args: + raw_output: Model output + strict_numbering: Only treat a leading number as list numbering when it is + followed by a delimiter and whitespace ("1. term"), so keywords such as + "2024" or "401k" are kept intact + Returns: List of cleaned keyword strings """ - import re - if not raw_output or not raw_output.strip(): return [] @@ -268,7 +571,12 @@ def _parse_keywords(self, raw_output: str) -> List[str]: line = re.sub(r"^[\-\*•▪→►]\s*", "", line) # Remove numbered prefixes: 1., 1), (1), etc. - line = re.sub(r"^[\(\[]?\d+[\)\]\.:\-]?\s*", "", line) + numbering = ( + r"^[\(\[]?\d+[\)\]\.:\-]\s+" + if strict_numbering + else r"^[\(\[]?\d+[\)\]\.:\-]?\s*" + ) + line = re.sub(numbering, "", line) # If line contains commas, it might be multiple keywords if "," in line: @@ -365,8 +673,12 @@ def reformulate(self, q: QueryItem, contexts=None) -> ReformulationResult: # Generate keywords out = self.llm.chat(msgs, temperature=temperature, max_tokens=max_tokens) - # Parse keywords using robust parser - terms = self._parse_keywords(out) + # Parse keywords, removing chat-model framing unless disabled + if self._clean_output_enabled(): + terms, cleanup_meta = self._extract_keywords(out) + metadata.update({"clean_output": True, **cleanup_meta}) + else: + terms = self._parse_keywords(out) # Limit to max 20 keywords max_keywords = int(self.cfg.params.get("max_keywords", 20)) diff --git a/reproducibility/tests/test_repro_schema.py b/reproducibility/tests/test_repro_schema.py index 0fc42cc..2baa608 100644 --- a/reproducibility/tests/test_repro_schema.py +++ b/reproducibility/tests/test_repro_schema.py @@ -743,3 +743,63 @@ def test_validator_rejects_absolute_topics(): payload["config"]["dataset_config"]["topics"] = "/abs/path/to/topics.tsv" with pytest.raises(ValidationError, match="absolute path"): validate(payload, skip_registry_checks=True) + + +def test_run_summary_keeps_empty_recorded_method_params(): + import importlib + + rs = importlib.import_module("examples.querygym_pyserini.run_summary") + results = { + "reformulation": { + "dataset": {"topics": "dl19-passage", "index": "msmarco-v1-passage", + "num_queries": 43, "bm25_weights": {"k1": 0.9, "b": 0.4}}, + "reformulation": {"method_params": {}, + "llm_config": {"temperature": 1.0, "max_tokens": 128}}, + "timing": {"total_time_seconds": 1.0}, + }, + "evaluation": {"timing": {"eval_time_seconds": 1.0}, "results": {"ndcg_cut_10": 0.5}}, + } + payload = rs._build_v1_summary( + results=results, + dataset_name="msmarco-v1-passage.trecdl2019", + method="query2e", + model="openai/gpt-4.1", + method_params={"clean_output": False}, + llm_config={"temperature": 1.0, "max_tokens": 128}, + steps=["reformulate", "retrieve", "evaluate"], + pipeline_time=3.0, + registry_path="dataset_registry.yaml", + queries_file=None, + index_name="msmarco-v1-passage", + ) + assert payload["config"]["method_params"] == {} + + +@pytest.mark.parametrize( + "run_path", + [ + "bright-biology/query2e/openai/gpt-4.1/bm25/edce8c2a.json", + "beir-v1.0.0-arguana/query2e/Qwen/Qwen2.5-72B-Instruct/bge-base-en-v1.5/eaea15f7.json", + ], +) +def test_query2e_clean_output_params_hash(run_path): + from querygym.core.base import MethodConfig + from querygym.core.prompts import PromptBank + from querygym.methods.query2e import Query2E + + root = Path(__file__).resolve().parents[2] + run = json.loads((root / "reproducibility" / "data" / "runs" / run_path).read_text()) + model = run["pipeline"]["model"] + params = run["config"]["method_params"] + llm_config = run["config"]["llm_config"] + bank = PromptBank(root / "querygym" / "prompt_bank.yaml") + + def params_hash(extra): + cfg = MethodConfig(name="query2e", params={**params, **extra}, llm={}) + effective = Query2E(cfg, None, bank).effective_params() + return compute_params_hash("query2e", model, effective, llm_config) + + # Raw-output parsing reproduces the recorded run identity + assert params_hash({"clean_output": False}) == run["params_hash"] + # Cleaned output (the default) is a different configuration + assert params_hash({}) != run["params_hash"] diff --git a/tests/test_methods_query2e.py b/tests/test_methods_query2e.py new file mode 100644 index 0000000..9dc8f47 --- /dev/null +++ b/tests/test_methods_query2e.py @@ -0,0 +1,211 @@ +from querygym.core.base import MethodConfig, QueryItem +from querygym.core.prompts import PromptBank +from querygym.core.utils import strip_think_trace +from querygym.methods.query2e import Query2E +from pathlib import Path +import pytest + +PROMPT_BANK = Path(__file__).parents[1] / "querygym" / "prompt_bank.yaml" +QUERY = "what is prime rate in canada" + +CATEGORIZED = """Here is a list of relevant keywords and phrases for the query "what is prime rate in canada": + +**Core Concepts** +* Prime Rate +* Prime Lending Rate + +**Canadian Context** +* Bank of Canada +* Overnight rate +""" + + +class DummyLLM: + def __init__(self, response): + self.response = response + + def chat(self, messages, **kwargs): + return self.response + + +def make_q2e(response="", **params): + cfg = MethodConfig(name="query2e", params=params, llm={"model": "dummy"}) + return Query2E(cfg, DummyLLM(response), PromptBank(PROMPT_BANK)) + + +def test_default_removes_chat_framing(): + res = make_q2e(CATEGORIZED).reformulate(QueryItem("q1", QUERY)) + + assert res.metadata["keywords"] == [ + "Prime Rate", + "Prime Lending Rate", + "Bank of Canada", + "Overnight rate", + ] + assert res.metadata["clean_output"] is True + assert "Here is" not in res.reformulated + + +def test_clean_output_false_parses_raw_output(): + meth = make_q2e(CATEGORIZED, clean_output=False) + res = meth.reformulate(QueryItem("q1", QUERY)) + + terms = meth._parse_keywords(CATEGORIZED) + assert res.metadata == { + "mode": "zs", + "prompt_id": "q2e.zs.v1", + "keywords": terms, + "temperature": 0.3, + "max_tokens": 256, + } + assert res.reformulated == meth.concatenate_result(QUERY, " ".join(terms)) + + +@pytest.mark.parametrize( + "raw, expected", + [ + ( + 'Sure! Here is a list of keywords for the query **"what is prime rate in canada":**\n\n' + "- prime rate\n- Bank of Canada\n\nLet me know if you need a more targeted list!", + ["prime rate", "Bank of Canada"], + ), + ( + "### Core Medical Procedures\n- Sclerotherapy\n- Endovenous laser ablation (EVLA)", + ["Sclerotherapy", "Endovenous laser ablation (EVLA)"], + ), + ( + "**Primary keywords:** prime rate, Canada\n**Related terms:** overnight rate", + ["prime rate", "Canada", "overnight rate"], + ), + ( + "1. **Prime rate** – the base lending rate\n2. **Bank of Canada** – central bank", + ["Prime rate", "Bank of Canada"], + ), + ( + "Okay, here's a breakdown:\n\n1. Core Concepts\n - Prime rate\n - Lending rate", + ["Prime rate", "Lending rate"], + ), + ( + "For this query, useful keywords include prime rate, Bank of Canada, and lending rate.", + ["prime rate", "Bank of Canada", "lending rate"], + ), + ( + "Keywords: prime rate, Bank of Canada, lending rate", + ["prime rate", "Bank of Canada", "lending rate"], + ), + ( + "| Category | Keywords |\n|---|---|\n| Core | prime rate, lending rate |", + ["prime rate", "lending rate"], + ), + ('```json\n["prime rate", "bank of canada"]\n```', ["prime rate", "bank of canada"]), + ( + "prime rate, Bank of Canada\n\nWould you like me to add related terms?", + ["prime rate", "Bank of Canada"], + ), + ( + "okay google, voice assistant, smart speaker", + ["okay google", "voice assistant", "smart speaker"], + ), + ("- 2024\n- 401k\n- 5-year fixed", ["2024", "401k", "5-year fixed"]), + ( + "Here is a list of keywords you can use: prime rate, bank of canada, lending rate, " + "overnight rate, mortgage rate, interest rate.", + [ + "prime rate", + "bank of canada", + "lending rate", + "overnight rate", + "mortgage rate", + "interest rate", + ], + ), + ("| Keyword | Why |\n|---|---|\n| prime rate | the base rate |", ["prime rate"]), + ( + "
      \n
    • prime rate
      • \n
    • bank of canada
      • \n
    ", + ["prime rate", "bank of canada"], + ), + ( + "Okay, the user asks about rates.\n    \n\nprime rate, bank of canada", + ["prime rate", "bank of canada"], + ), + ("- __init__ method\n- snake_case", ["__init__ method", "snake_case"]), + ], +) +def test_cleanup_keeps_only_keywords(raw, expected): + terms, _ = make_q2e()._extract_keywords(raw) + assert terms == expected + + +def test_think_trace_is_removed(): + terms, meta = make_q2e()._extract_keywords( + "\nThe user wants keywords.\n\n\nprime rate, Bank of Canada" + ) + assert terms == ["prime rate", "Bank of Canada"] + assert meta == {} + + +def test_unclosed_think_trace_yields_no_keywords(): + terms, meta = make_q2e()._extract_keywords( + "\nOkay, the user is asking about the prime rate. Related terms: lending," + ) + assert terms == [] + assert meta == {"unclosed_think": True} + + +def test_few_shot_continuation_is_truncated(): + terms, meta = make_q2e()._extract_keywords( + "prime rate, bank of canada\nQuery: how tall is mount everest\nKeywords: everest, nepal" + ) + assert terms == ["prime rate", "bank of canada"] + assert meta == {"truncated_at_query": True} + + +@pytest.mark.parametrize( + "raw", + [ + "Here are the keywords for your query.", + "\nkeywords: prime rate, bank of canada\n\n", + ], +) +def test_output_without_keywords_is_flagged(raw): + terms, meta = make_q2e()._extract_keywords(raw) + assert terms == [] + assert meta == {"cleanup_empty": True} + + +def test_cleanup_falls_back_to_plain_keyword_list(): + terms, meta = make_q2e()._extract_keywords("note, musical notation, pitch") + assert terms == ["note", "musical notation", "pitch"] + assert meta == {"cleanup_fallback": True} + + +@pytest.mark.parametrize("raw", [None, "", " \n"]) +def test_empty_output_expands_with_query_only(raw): + res = make_q2e(raw).reformulate(QueryItem("q1", QUERY)) + assert res.metadata["keywords"] == [] + assert "error" not in res.metadata + + +@pytest.mark.parametrize( + "params, expected", + [ + ({}, {"clean_output": True}), + ({"clean_output": True, "mode": "zs"}, {"clean_output": True, "mode": "zs"}), + ({"clean_output": False}, {}), + ({"clean_output": False, "num_examples": 4}, {"num_examples": 4}), + ], +) +def test_effective_params_records_clean_output_only_when_enabled(params, expected): + assert make_q2e(**params).effective_params() == expected + + +def test_effective_params_drops_runtime_objects(): + meth = make_q2e(clean_output=False, searcher=object(), retrieval_k=10) + assert meth.effective_params() == {"retrieval_k": 10} + + +def test_strip_think_trace(): + assert strip_think_trace("ab") == ("b", False) + assert strip_think_trace("xunfinished") == ("x", True) + assert strip_think_trace("xy") == ("y", False) + assert strip_think_trace("plain") == ("plain", False)