query2e: remove chat-model framing from LLM output

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 `<think>` reasoning traces
+- `--clean-output/--no-clean-output` CLI flag
+
+### Changed
+- Query2E removes chat-model framing (`<think>` 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

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

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 `qid<TAB>query`.
 - `prompts-list` / `prompts-show` — inspect prompt bank.
 - `script-gen` — create a Pyserini + trec_eval bash script.

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: `<think>` 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)
 
 ---
 

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, `<think>` 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")

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)

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:

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)))

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:

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:

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"<think>.*?</think>", re.DOTALL)
+
+
+def strip_think_trace(text: str) -> Tuple[str, bool]:
+    """Remove <think>...</think> reasoning traces from model output.
+
+    Args:
+        text: Raw model output
+
+    Returns:
+        (text without reasoning traces, whether an unclosed <think> was found).
+        When a <think> 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 <think> in the prompt),
+        everything up to it is dropped.
+    """
+    text = _THINK_BLOCK.sub("", text)
+    if "</think>" in text:
+        text = text.split("</think>", 1)[1]
+    if "<think>" in text:
+        return text.split("<think>", 1)[0], True
+    return text, False
 
 
 def seed_everything(seed: Optional[int] = 42):