├── examples
    ├── docker
    │   ├── docker-compose.yml
    │   ├── README.md
    │   └── quickstart.ipynb
    ├── snippets
    │   ├── tiny_qrels.txt
    │   ├── tiny_queries.tsv
    │   ├── tiny_contexts.jsonl
    │   ├── 02_load_with_loaders.py
    │   ├── 03_query_reformulation.py
    │   ├── 01_load_from_file.py
    │   ├── 05_using_adapters.py
    │   ├── 04_context_reformulation.py
    │   └── README.md
    ├── README.md
    └── querygym_pyserini
    │   └── reformulation_config.yaml
├── .env.example
├── docs
    ├── querygym-logo.gif
    ├── querygym-logo.png
    ├── user-guide
    │   ├── cli.md
    │   ├── prompts.md
    │   ├── datasets.md
    │   └── reformulation.md
    ├── development
    │   ├── guide.md
    │   └── contributing.md
    ├── about
    │   ├── license.md
    │   └── changelog.md
    ├── api
    │   ├── methods.md
    │   ├── adapters.md
    │   └── loaders.md
    ├── getting-started
    │   ├── installation.md
    │   └── quickstart.md
    └── index.md
├── index.md
├── scripts
    └── templated_retrieval.sh
├── Gemfile
├── querygym
    ├── __main__.py
    ├── methods
    │   ├── __init__.py
    │   ├── genqr.py
    │   ├── lamer.py
    │   ├── mugi.py
    │   ├── genqr_ensemble.py
    │   ├── qa_expand.py
    │   └── csqe.py
    ├── core
    │   ├── registry.py
    │   ├── utils.py
    │   ├── runner.py
    │   ├── llm.py
    │   ├── searcher.py
    │   └── prompts.py
    ├── loaders
    │   ├── __init__.py
    │   ├── msmarco.py
    │   └── beir.py
    ├── adapters
    │   ├── __init__.py
    │   ├── pyterrier_adapter.py
    │   └── pyserini_adapter.py
    ├── config
    │   └── defaults.yaml
    └── __init__.py
├── tests
    ├── test_prompts.py
    ├── test_cli_script_gen.py
    ├── test_dataloader.py
    └── test_methods_genqr.py
├── .readthedocs.yaml
├── .github
    ├── ISSUE_TEMPLATE
    │   ├── question.md
    │   ├── feature_request.md
    │   └── bug_report.md
    ├── workflows
    │   ├── publish.yml
    │   ├── pages.yml
    │   └── docker.yml
    ├── PULL_REQUEST_TEMPLATE.md
    └── CONTRIBUTING.md
├── .dockerignore
├── MANIFEST.in
├── _config.yml
├── DOCKER_SETUP.md
├── Dockerfile
├── docker-compose.yml
├── Makefile
├── Dockerfile.cpu
├── .gitignore
├── mkdocs.yml
├── pyproject.toml
├── README.md
└── _layouts
    └── default.html


/examples/docker/docker-compose.yml:
--------------------------------------------------------------------------------
1 | ../../docker-compose.yml


--------------------------------------------------------------------------------
/examples/snippets/tiny_qrels.txt:
--------------------------------------------------------------------------------
1 | 1	0	doc1	1
2 | 1	0	doc2	2
3 | 2	0	doc3	1
4 | 


--------------------------------------------------------------------------------
/.env.example:
--------------------------------------------------------------------------------
1 | # Environment variables example
2 | # OPENAI_API_KEY=your_api_key_here
3 | 


--------------------------------------------------------------------------------
/examples/snippets/tiny_queries.tsv:
--------------------------------------------------------------------------------
1 | 1	what is neural ir
2 | 2	bert reranking ms marco
3 | 


--------------------------------------------------------------------------------
/docs/querygym-logo.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ls3-lab/QueryGym/HEAD/docs/querygym-logo.gif


--------------------------------------------------------------------------------
/docs/querygym-logo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ls3-lab/QueryGym/HEAD/docs/querygym-logo.png


--------------------------------------------------------------------------------
/index.md:
--------------------------------------------------------------------------------
1 | ---
2 | layout: default
3 | title: Home
4 | ---
5 | 
6 | {% include_relative README.md %}
7 | 


--------------------------------------------------------------------------------
/scripts/templated_retrieval.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 | set -euo pipefail
3 | echo "This is an example produced by 'querygym script-gen'."
4 | 


--------------------------------------------------------------------------------
/Gemfile:
--------------------------------------------------------------------------------
1 | source "https://rubygems.org"
2 | 
3 | gem "jekyll", "~> 4.3"
4 | gem "kramdown-parser-gfm" # GitHub Flavored Markdown
5 | gem "webrick" # Required for Ruby 3.0+
6 | 


--------------------------------------------------------------------------------
/querygym/__main__.py:
--------------------------------------------------------------------------------
1 | """Entry point for running queryGym as a module: python -m queryGym"""
2 | from .cli import app
3 | 
4 | if __name__ == "__main__":
5 |     app()
6 | 
7 | 


--------------------------------------------------------------------------------
/tests/test_prompts.py:
--------------------------------------------------------------------------------
1 | from querygym.core.prompts import PromptBank
2 | from pathlib import Path
3 | 
4 | def test_prompt_bank_loads():
5 |     pb = PromptBank(Path(__file__).parents[1] / "querygym" / "prompt_bank.yaml")
6 |     assert len(pb.list()) > 0
7 | 


--------------------------------------------------------------------------------
/docs/user-guide/cli.md:
--------------------------------------------------------------------------------
1 | # CLI
2 | 
3 | - `run` — produce reformulated TSV.
4 | - `data-to-tsv` — convert any backend to `qid<TAB>query`.
5 | - `prompts-list` / `prompts-show` — inspect prompt bank.
6 | - `script-gen` — create a Pyserini + trec_eval bash script.
7 | 


--------------------------------------------------------------------------------
/tests/test_cli_script_gen.py:
--------------------------------------------------------------------------------
1 | from querygym.cli import build_script_lines
2 | 
3 | def test_script_lines():
4 |     lines = build_script_lines(index_path="/idx", topics="/q.tsv", run="/run.txt", qrels="/qrels.txt")
5 |     assert any("--index /idx" in ln for ln in lines)
6 | 


--------------------------------------------------------------------------------
/querygym/methods/__init__.py:
--------------------------------------------------------------------------------
1 | from .genqr import GENQR
2 | from .genqr_ensemble import GenQREnsemble
3 | from .query2doc import Query2Doc
4 | from .qa_expand import QAExpand
5 | from .mugi import MuGI
6 | from .lamer import LameR
7 | from .query2e import Query2E
8 | from .csqe import CSQE
9 | 


--------------------------------------------------------------------------------
/examples/snippets/tiny_contexts.jsonl:
--------------------------------------------------------------------------------
1 | {"qid": "1", "contexts": ["Neural information retrieval uses deep learning models for search", "BERT and transformers revolutionized IR systems"]}
2 | {"qid": "2", "contexts": ["BERT reranking improves MS MARCO passage ranking", "Cross-encoders are effective for reranking tasks"]}
3 | 


--------------------------------------------------------------------------------
/querygym/core/registry.py:
--------------------------------------------------------------------------------
 1 | from __future__ import annotations
 2 | from typing import Dict, Type
 3 | from .base import BaseReformulator
 4 | 
 5 | METHODS: Dict[str, Type[BaseReformulator]] = {}
 6 | 
 7 | def register_method(name: str):
 8 |     def deco(cls: Type[BaseReformulator]):
 9 |         METHODS[name] = cls
10 |         return cls
11 |     return deco
12 | 


--------------------------------------------------------------------------------
/docs/user-guide/prompts.md:
--------------------------------------------------------------------------------
 1 | # Prompt Bank
 2 | 
 3 | Single YAML file: `queryGym/prompt_bank.yaml`
 4 | 
 5 | Each entry:
 6 | - `id`: Unique string (e.g., `genqr.keywords.v1`)
 7 | - `method_family`: logical grouping
 8 | - `version`: integer
 9 | - `introduced_by`: citation-like text
10 | - `license`, `authors`, `tags`
11 | - `template.system`, `template.user`: may contain `{query}`
12 | - `notes`: free text
13 | 


--------------------------------------------------------------------------------
/.readthedocs.yaml:
--------------------------------------------------------------------------------
 1 | # Read the Docs configuration file
 2 | # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
 3 | 
 4 | version: 2
 5 | 
 6 | build:
 7 |   os: ubuntu-24.04
 8 |   tools:
 9 |     python: "3.9"
10 | 
11 | mkdocs:
12 |   configuration: mkdocs.yml
13 | 
14 | python:
15 |   install:
16 |     - method: pip
17 |       path: .
18 |       extra_requirements:
19 |         - dev
20 |         


--------------------------------------------------------------------------------
/tests/test_dataloader.py:
--------------------------------------------------------------------------------
 1 | from querygym.data.dataloader import DataLoader
 2 | from pathlib import Path
 3 | import tempfile
 4 | 
 5 | def test_local_tsv_loading():
 6 |     with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".tsv") as f:
 7 |         f.write("1\tapple\n2\tbanana\n")
 8 |         path = f.name
 9 |     queries = DataLoader.load_queries(path, format="tsv")
10 |     assert len(queries) == 2 and queries[0].qid == "1"
11 | 


--------------------------------------------------------------------------------
/querygym/loaders/__init__.py:
--------------------------------------------------------------------------------
 1 | """Format-specific loaders for queryGym.
 2 | 
 3 | This module provides convenience functions for loading data from popular IR dataset
 4 | formats like BEIR and MS MARCO. These loaders assume you've already downloaded the
 5 | data using the respective dataset libraries.
 6 | 
 7 | For general file loading, use queryGym.data.DataLoader instead.
 8 | """
 9 | 
10 | from . import beir, msmarco
11 | 
12 | __all__ = ["beir", "msmarco"]
13 | 


--------------------------------------------------------------------------------
/querygym/adapters/__init__.py:
--------------------------------------------------------------------------------
 1 | """
 2 | Searcher Adapters for queryGym
 3 | 
 4 | This package contains adapters that implement the BaseSearcher interface
 5 | for different search libraries like Pyserini and PyTerrier.
 6 | """
 7 | 
 8 | # Import adapters to register them
 9 | from .pyserini_adapter import PyseriniSearcher
10 | from .pyterrier_adapter import PyTerrierSearcher
11 | 
12 | __all__ = [
13 |     "PyseriniSearcher",
14 |     "PyTerrierSearcher",
15 | ]
16 | 
17 | 


--------------------------------------------------------------------------------
/docs/user-guide/datasets.md:
--------------------------------------------------------------------------------
 1 | # Datasets
 2 | 
 3 | Use **UnifiedQuerySource** to load Local / MS MARCO / BEIR.
 4 | 
 5 | - Local TSV: `qid<TAB>query`
 6 | - Local JSONL: `{"qid": "...", "query": "..."}`
 7 | 
 8 | MS MARCO:
 9 | - `backend=msmarco`, `source=file`, `--msmarco-queries-tsv ...`
10 | - or `source=hf`, `--hf-name ms_marco`, `--hf-config passage|document`, `--split dev|train|...`
11 | 
12 | BEIR:
13 | - `backend=beir`, `source=beir`, `--beir-root /path/to/fiqa` (expects BEIR layout)
14 | - or `source=hf`, `--hf-name beir/fiqa`, `--split test`
15 | 


--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/question.md:
--------------------------------------------------------------------------------
 1 | ---
 2 | name: Question
 3 | about: Ask a question about using QueryGym
 4 | title: '[QUESTION] '
 5 | labels: question
 6 | assignees: ''
 7 | ---
 8 | 
 9 | ## Question
10 | 
11 | What would you like to know?
12 | 
13 | ## Context
14 | 
15 | Provide any relevant context:
16 | - What are you trying to achieve?
17 | - What have you tried so far?
18 | - What documentation have you consulted?
19 | 
20 | ## Code (if applicable)
21 | 
22 | ```python
23 | # Relevant code snippet
24 | ```
25 | 
26 | ## Environment (if relevant)
27 | 
28 | - **QueryGym version**: 
29 | - **Python version**: 
30 | 


--------------------------------------------------------------------------------
/querygym/core/utils.py:
--------------------------------------------------------------------------------
 1 | from __future__ import annotations
 2 | import random
 3 | from typing import Optional
 4 | 
 5 | def seed_everything(seed: Optional[int] = 42):
 6 |     if seed is None: return
 7 |     random.seed(seed)
 8 |     try:
 9 |         import numpy as np
10 |         np.random.seed(seed)
11 |     except Exception:
12 |         pass
13 |     try:
14 |         import torch
15 |         torch.manual_seed(seed)
16 |         torch.cuda.manual_seed_all(seed)
17 |         torch.backends.cudnn.deterministic = True
18 |         torch.backends.cudnn.benchmark = False
19 |     except Exception:
20 |         pass
21 | 


--------------------------------------------------------------------------------
/examples/snippets/02_load_with_loaders.py:
--------------------------------------------------------------------------------
 1 | """Load BEIR and MS MARCO datasets"""
 2 | import querygym as qg
 3 | 
 4 | # BEIR dataset (assumes already downloaded)
 5 | beir_path = "./data/beir/nfcorpus"
 6 | queries = qg.loaders.beir.load_queries(beir_path, split="test")
 7 | qrels = qg.loaders.beir.load_qrels(beir_path, split="test")
 8 | print(f"BEIR: {len(queries)} queries, {len(qrels)} qrels")
 9 | 
10 | # MS MARCO dataset
11 | queries = qg.loaders.msmarco.load_queries("./data/msmarco/queries.tsv")
12 | qrels = qg.loaders.msmarco.load_qrels("./data/msmarco/qrels.tsv")
13 | print(f"MS MARCO: {len(queries)} queries, {len(qrels)} qrels")
14 | 


--------------------------------------------------------------------------------
/examples/snippets/03_query_reformulation.py:
--------------------------------------------------------------------------------
 1 | """Reformulate queries with different methods"""
 2 | import querygym as qg
 3 | from pathlib import Path
 4 | 
 5 | example_dir = Path(__file__).parent
 6 | 
 7 | # Load queries
 8 | queries = qg.load_queries(example_dir / "tiny_queries.tsv")
 9 | 
10 | # Create reformulator
11 | reformulator = qg.create_reformulator("genqr_ensemble", model="gpt-4")
12 | 
13 | # Reformulate
14 | results = reformulator.reformulate_batch(queries)
15 | 
16 | # Show results
17 | for r in results:
18 |     print(f"{r.qid}:")
19 |     print(f"  Original: {r.original}")
20 |     print(f"  Reformed: {r.reformulated}\n")
21 | 


--------------------------------------------------------------------------------
/examples/snippets/01_load_from_file.py:
--------------------------------------------------------------------------------
 1 | """Load queries from local files"""
 2 | import querygym as qg
 3 | from pathlib import Path
 4 | 
 5 | example_dir = Path(__file__).parent
 6 | 
 7 | # Load from TSV
 8 | queries = qg.load_queries(example_dir / "tiny_queries.tsv")
 9 | print(f"Loaded {len(queries)} queries")
10 | for q in queries[:2]:
11 |     print(f"  {q.qid}: {q.text}")
12 | 
13 | # Load qrels
14 | qrels = qg.load_qrels(example_dir / "tiny_qrels.txt")
15 | print(f"\nLoaded qrels for {len(qrels)} queries")
16 | 
17 | # Save queries
18 | qg.DataLoader.save_queries(queries, example_dir / "output.tsv", format="tsv")
19 | print("Saved to output.tsv")
20 | 


--------------------------------------------------------------------------------
/examples/snippets/05_using_adapters.py:
--------------------------------------------------------------------------------
 1 | """Use Pyserini and PyTerrier adapters"""
 2 | import querygym as qg
 3 | 
 4 | # Pyserini adapter
 5 | from pyserini.search.lucene import LuceneSearcher
 6 | searcher = LuceneSearcher.from_prebuilt_index("msmarco-v1-passage")
 7 | pyserini_adapter = qg.wrap_pyserini_searcher(searcher)
 8 | 
 9 | # Search
10 | query = qg.QueryItem("1", "what is neural IR")
11 | hits = pyserini_adapter.search(query.text, k=10)
12 | print(f"Pyserini: {len(hits)} hits")
13 | 
14 | # PyTerrier adapter (if available)
15 | # import pyterrier as pt
16 | # bm25 = pt.BatchRetrieve.from_dataset("msmarco_passage", "terrier_stemmed")
17 | # pyterrier_adapter = qg.wrap_pyterrier_retriever(bm25)
18 | 


--------------------------------------------------------------------------------
/examples/snippets/04_context_reformulation.py:
--------------------------------------------------------------------------------
 1 | """Reformulate queries using retrieved contexts"""
 2 | import querygym as qg
 3 | from pathlib import Path
 4 | 
 5 | example_dir = Path(__file__).parent
 6 | 
 7 | # Load queries and contexts
 8 | queries = qg.load_queries(example_dir / "tiny_queries.tsv")
 9 | contexts = qg.load_contexts(example_dir / "tiny_contexts.jsonl")
10 | 
11 | # Create context-based reformulator (Query2Doc uses contexts)
12 | reformulator = qg.create_reformulator("query2doc", model="gpt-4")
13 | 
14 | # Reformulate with contexts
15 | results = reformulator.reformulate_batch(queries, contexts=contexts)
16 | 
17 | # Show results
18 | for r in results:
19 |     print(f"{r.qid}: {r.reformulated[:80]}...")
20 | 


--------------------------------------------------------------------------------
/.dockerignore:
--------------------------------------------------------------------------------
 1 | # Python
 2 | __pycache__/
 3 | *.py[cod]
 4 | *$py.class
 5 | *.so
 6 | .Python
 7 | *.egg-info/
 8 | dist/
 9 | build/
10 | *.egg
11 | 
12 | # Virtual environments
13 | venv/
14 | env/
15 | ENV/
16 | 
17 | # IDE
18 | .vscode/
19 | .idea/
20 | *.swp
21 | *.swo
22 | *~
23 | 
24 | # Git
25 | .git/
26 | .gitignore
27 | 
28 | # Documentation build
29 | site/
30 | docs/_build/
31 | 
32 | # Testing
33 | .pytest_cache/
34 | .coverage
35 | htmlcov/
36 | .tox/
37 | 
38 | # Data and outputs (will be mounted as volumes)
39 | data/
40 | outputs/
41 | notebooks/
42 | *.tsv
43 | *.jsonl
44 | *.json
45 | *.txt
46 | 
47 | # OS
48 | .DS_Store
49 | Thumbs.db
50 | 
51 | # Docker
52 | Dockerfile*
53 | docker-compose.yml
54 | .dockerignore
55 | 


--------------------------------------------------------------------------------
/MANIFEST.in:
--------------------------------------------------------------------------------
 1 | # Include package data files
 2 | include queryGym/prompt_bank.yaml
 3 | include queryGym/config/*.yaml
 4 | 
 5 | # Include documentation
 6 | include README.md
 7 | include LICENSE
 8 | include CONTRIBUTING.md
 9 | include Structure.md
10 | 
11 | # Include examples
12 | recursive-include examples *.py *.tsv *.ipynb
13 | 
14 | # Include docs
15 | recursive-include docs *.md
16 | 
17 | # Exclude development and build artifacts
18 | global-exclude __pycache__
19 | global-exclude *.py[co]
20 | global-exclude .DS_Store
21 | global-exclude *.egg-info
22 | global-exclude *.so
23 | global-exclude .pytest_cache
24 | global-exclude htmlcov
25 | global-exclude .coverage
26 | global-exclude .mypy_cache
27 | global-exclude .ruff_cache
28 | 


--------------------------------------------------------------------------------
/tests/test_methods_genqr.py:
--------------------------------------------------------------------------------
 1 | from querygym.core.base import MethodConfig, QueryItem
 2 | from querygym.core.prompts import PromptBank
 3 | from querygym.methods.genqr_ensemble import GenQREnsemble
 4 | from pathlib import Path
 5 | 
 6 | class DummyLLM:
 7 |     def chat(self, messages, **kwargs):
 8 |         return "bertopic"
 9 | 
10 | def test_genqr_ensemble():
11 |     cfg = MethodConfig(name="genqr_ensemble", params={"repeat_query_weight":2}, llm={"model":"dummy"})
12 |     llm = DummyLLM()
13 |     pb = PromptBank(Path(__file__).parents[1] / "querygym" / "prompt_bank.yaml")
14 |     meth = GenQREnsemble(cfg, llm, pb)
15 |     res = meth.reformulate(QueryItem("Q1", "bertopic"))
16 |     assert "bertopic bertopic" in res.reformulated
17 | 


--------------------------------------------------------------------------------
/_config.yml:
--------------------------------------------------------------------------------
 1 | # Jekyll configuration for GitHub Pages
 2 | title: QueryGym
 3 | description: A lightweight, reproducible toolkit for LLM-based query reformulation
 4 | url: "https://ls3-lab.github.io"
 5 | baseurl: "/QueryGym"
 6 | 
 7 | # Markdown settings
 8 | markdown: kramdown
 9 | highlighter: rouge
10 | 
11 | # Exclude files from build
12 | exclude:
13 |   - Gemfile
14 |   - Gemfile.lock
15 |   - node_modules
16 |   - vendor
17 |   - .git
18 |   - .github
19 |   - tests
20 |   - querygym
21 |   - examples
22 |   - scripts
23 |   - "*.py"
24 |   - "*.toml"
25 |   - Makefile
26 |   - Dockerfile
27 |   - Dockerfile.cpu
28 |   - docker-compose.yml
29 |   - .dockerignore
30 |   - .gitignore
31 |   - requirements*.txt
32 |   - docs/
33 | 
34 | # Include specific files
35 | include:
36 |   - README.md
37 |   - index.md
38 |   - LICENSE
39 | 


--------------------------------------------------------------------------------
/.github/workflows/publish.yml:
--------------------------------------------------------------------------------
 1 | name: Publish to PyPI
 2 | 
 3 | on:
 4 |   release:
 5 |     types: [published]
 6 | 
 7 | jobs:
 8 |   deploy:
 9 |     runs-on: ubuntu-latest
10 |     
11 |     steps:
12 |     - uses: actions/checkout@v4
13 |     
14 |     - name: Set up Python
15 |       uses: actions/setup-python@v5
16 |       with:
17 |         python-version: '3.9'
18 |     
19 |     - name: Install dependencies
20 |       run: |
21 |         python -m pip install --upgrade pip
22 |         pip install build twine
23 |     
24 |     - name: Build package
25 |       run: |
26 |         rm -rf dist build *.egg-info
27 |         python -m build
28 |     
29 |     - name: Check package
30 |       run: twine check dist/*
31 |     
32 |     - name: Publish to PyPI
33 |       env:
34 |         TWINE_USERNAME: __token__
35 |         TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
36 |       run: twine upload --skip-existing dist/*
37 | 


--------------------------------------------------------------------------------
/examples/snippets/README.md:
--------------------------------------------------------------------------------
 1 | # queryGym Code Snippets
 2 | 
 3 | Quick, focused examples showing core functionality.
 4 | 
 5 | ## Examples
 6 | 
 7 | 1. **`01_load_from_file.py`** - Load queries and qrels from local files
 8 | 2. **`02_load_with_loaders.py`** - Load BEIR and MS MARCO datasets
 9 | 3. **`03_query_reformulation.py`** - Reformulate queries with LLMs
10 | 4. **`04_context_reformulation.py`** - Reformulate using retrieved contexts
11 | 5. **`05_using_adapters.py`** - Use Pyserini and PyTerrier adapters
12 | 
13 | ## Run Examples
14 | 
15 | ```bash
16 | cd examples
17 | conda activate querygym
18 | 
19 | # Set API key for reformulation
20 | export OPENAI_API_KEY='sk-...'
21 | 
22 | # Run any example
23 | python 01_load_from_file.py
24 | python 03_query_reformulation.py
25 | ```
26 | 
27 | ## Data Files
28 | 
29 | Examples use:
30 | - `tiny_queries.tsv` - Sample queries
31 | - `tiny_qrels.txt` - Sample relevance judgments
32 | - `tiny_contexts.jsonl` - Sample retrieved contexts
33 | - BEIR/MS MARCO data (download separately)
34 | 


--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/feature_request.md:
--------------------------------------------------------------------------------
 1 | ---
 2 | name: Feature Request
 3 | about: Suggest a new feature or improvement for QueryGym
 4 | title: '[FEATURE] '
 5 | labels: enhancement
 6 | assignees: ''
 7 | ---
 8 | 
 9 | ## Feature Description
10 | 
11 | A clear and concise description of the feature you'd like.
12 | 
13 | ## Use Case
14 | 
15 | Describe the problem this feature would solve or the use case it enables.
16 | 
17 | ## Proposed Solution
18 | 
19 | How do you envision this feature working?
20 | 
21 | ## Example Usage
22 | 
23 | ```python
24 | # How would you use this feature?
25 | import querygym as qg
26 | 
27 | # Example code
28 | ```
29 | 
30 | ## Alternatives Considered
31 | 
32 | Have you considered any alternative solutions or workarounds?
33 | 
34 | ## Additional Context
35 | 
36 | Add any other context, mockups, or examples about the feature request here.
37 | 
38 | ## Checklist
39 | 
40 | - [ ] I have searched existing issues to ensure this is not a duplicate
41 | - [ ] I have considered if this fits the scope of QueryGym
42 | 


--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/bug_report.md:
--------------------------------------------------------------------------------
 1 | ---
 2 | name: Bug Report
 3 | about: Report a bug to help us improve QueryGym
 4 | title: '[BUG] '
 5 | labels: bug
 6 | assignees: ''
 7 | ---
 8 | 
 9 | ## Bug Description
10 | 
11 | A clear and concise description of the bug.
12 | 
13 | ## Steps to Reproduce
14 | 
15 | 1. 
16 | 2. 
17 | 3. 
18 | 
19 | ## Expected Behavior
20 | 
21 | What you expected to happen.
22 | 
23 | ## Actual Behavior
24 | 
25 | What actually happened.
26 | 
27 | ## Code Example
28 | 
29 | ```python
30 | # Minimal code to reproduce the issue
31 | import querygym as qg
32 | 
33 | # Your code here
34 | ```
35 | 
36 | ## Error Message
37 | 
38 | ```
39 | Paste the full error message/traceback here
40 | ```
41 | 
42 | ## Environment
43 | 
44 | - **OS**: [e.g., Ubuntu 22.04, Windows 11, macOS 14]
45 | - **Python version**: [e.g., 3.10.12]
46 | - **QueryGym version**: [e.g., 0.1.4]
47 | - **Installation method**: [pip, docker, source]
48 | 
49 | ## Additional Context
50 | 
51 | Add any other context about the problem here (screenshots, logs, etc.).
52 | 


--------------------------------------------------------------------------------
/querygym/core/runner.py:
--------------------------------------------------------------------------------
 1 | from __future__ import annotations
 2 | from typing import List, Optional, Dict
 3 | from .base import QueryItem, MethodConfig
 4 | from .llm import OpenAICompatibleClient
 5 | from .prompts import PromptBank
 6 | from .registry import METHODS
 7 | 
 8 | # Import all methods to ensure they're registered
 9 | from ..methods.genqr import GENQR
10 | from ..methods.genqr_ensemble import GenQREnsemble
11 | from ..methods.query2doc import Query2Doc
12 | from ..methods.qa_expand import QAExpand
13 | from ..methods.mugi import MuGI
14 | from ..methods.lamer import LameR
15 | from ..methods.query2e import Query2E
16 | 
17 | def build_llm(cfg: MethodConfig):
18 |     llm_cfg = cfg.llm
19 |     return OpenAICompatibleClient(model=llm_cfg["model"],
20 |                                   base_url=llm_cfg.get("base_url"),
21 |                                   api_key=llm_cfg.get("api_key"))
22 | 
23 | def run_method(method_name: str, cfg: MethodConfig, queries: List[QueryItem],
24 |                prompt_bank_path: str, ctx_map: Optional[Dict[str, List[str]]] = None):
25 |     Method = METHODS[method_name]
26 |     llm = build_llm(cfg)
27 |     pb = PromptBank(prompt_bank_path)
28 |     reformulator = Method(cfg, llm, pb)
29 |     return reformulator.reformulate_batch(queries, ctx_map)
30 | 


--------------------------------------------------------------------------------
/.github/PULL_REQUEST_TEMPLATE.md:
--------------------------------------------------------------------------------
 1 | ## Description
 2 | 
 3 | Brief description of the changes in this PR.
 4 | 
 5 | ## Type of Change
 6 | 
 7 | - [ ] 🐛 Bug fix (non-breaking change that fixes an issue)
 8 | - [ ] ✨ New feature (non-breaking change that adds functionality)
 9 | - [ ] 💥 Breaking change (fix or feature that would cause existing functionality to change)
10 | - [ ] 📚 Documentation update
11 | - [ ] 🔧 Refactoring (no functional changes)
12 | - [ ] 🧪 Test update
13 | 
14 | ## Related Issues
15 | 
16 | Closes #(issue number)
17 | 
18 | ## Changes Made
19 | 
20 | - 
21 | - 
22 | - 
23 | 
24 | ## Testing
25 | 
26 | Describe how you tested your changes:
27 | 
28 | - [ ] Unit tests pass (`pytest`)
29 | - [ ] Manual testing performed
30 | - [ ] New tests added for new functionality
31 | 
32 | ## Checklist
33 | 
34 | - [ ] My code follows the project's style guidelines
35 | - [ ] I have performed a self-review of my code
36 | - [ ] I have commented my code, particularly in hard-to-understand areas
37 | - [ ] I have updated the documentation accordingly
38 | - [ ] My changes generate no new warnings
39 | - [ ] Any dependent changes have been merged and published
40 | 
41 | ## Screenshots (if applicable)
42 | 
43 | Add screenshots to help explain your changes.
44 | 
45 | ## Additional Notes
46 | 
47 | Any additional information reviewers should know.
48 | 


--------------------------------------------------------------------------------
/docs/development/guide.md:
--------------------------------------------------------------------------------
 1 | # Development Guide
 2 | 
 3 | ## Setup
 4 | 
 5 | ```bash
 6 | python -m venv venv
 7 | source venv/bin/activate
 8 | pip install -e ".[dev,all]"
 9 | ```
10 | 
11 | The `-e` flag installs in editable mode: code changes take effect immediately without reinstalling.
12 | 
13 | ## Testing
14 | 
15 | ```bash
16 | pytest                    # Run all tests
17 | pytest --cov=queryGym    # With coverage
18 | black queryGym/ tests/   # Format code
19 | ruff check --fix queryGym/  # Lint
20 | ```
21 | 
22 | ## Adding a New Method
23 | 
24 | 1. Create `queryGym/methods/my_method.py`:
25 | ```python
26 | from ..core.base import BaseReformulator, QueryItem, ReformulationResult
27 | from ..core.registry import register_method
28 | 
29 | @register_method("my_method")
30 | class MyMethod(BaseReformulator):
31 |     def reformulate(self, q: QueryItem, contexts=None):
32 |         msgs = self.prompts.render("my_method.v1", query=q.text)
33 |         result = self.llm.chat(msgs, temperature=0.7)
34 |         return ReformulationResult(q.qid, q.text, result)
35 | ```
36 | 
37 | 2. Add prompt to `prompt_bank.yaml`
38 | 3. Export in `methods/__init__.py` and main `__init__.py`
39 | 4. Write tests in `tests/test_my_method.py`
40 | 
41 | ## Publishing
42 | 
43 | ```bash
44 | python -m build
45 | twine upload --repository testpypi dist/*  # Test first
46 | twine upload dist/*                         # Then publish
47 | 


--------------------------------------------------------------------------------
/.github/workflows/pages.yml:
--------------------------------------------------------------------------------
 1 | # GitHub Pages deployment - Main website with README as homepage
 2 | name: Deploy to GitHub Pages
 3 | 
 4 | on:
 5 |   push:
 6 |     branches: ["main"]
 7 |     paths:
 8 |       - 'README.md'
 9 |       - 'index.md'
10 |       - 'docs/leaderboard.html'
11 |       - 'docs/querygym-logo.png'
12 |       - '_config.yml'
13 |       - '_layouts/**'
14 |       - 'Gemfile'
15 |       - '.github/workflows/pages.yml'
16 |   workflow_dispatch:
17 | 
18 | permissions:
19 |   contents: read
20 |   pages: write
21 |   id-token: write
22 | 
23 | concurrency:
24 |   group: "pages"
25 |   cancel-in-progress: false
26 | 
27 | jobs:
28 |   build:
29 |     runs-on: ubuntu-latest
30 |     steps:
31 |       - name: Checkout
32 |         uses: actions/checkout@v4
33 | 
34 |       - name: Setup Pages
35 |         uses: actions/configure-pages@v4
36 | 
37 |       - name: Setup Ruby
38 |         uses: ruby/setup-ruby@v1
39 |         with:
40 |           ruby-version: '3.2'
41 |           bundler-cache: true
42 | 
43 |       - name: Install dependencies
44 |         run: bundle install
45 | 
46 |       - name: Build with Jekyll
47 |         run: bundle exec jekyll build
48 |         env:
49 |           JEKYLL_ENV: production
50 | 
51 |       - name: Copy additional files
52 |         run: |
53 |           # Copy leaderboard to _site
54 |           cp docs/leaderboard.html _site/leaderboard.html
55 |           # Copy logo for the site
56 |           cp docs/querygym-logo.png _site/querygym-logo.png 2>/dev/null || true
57 | 
58 |       - name: Upload artifact
59 |         uses: actions/upload-pages-artifact@v3
60 | 
61 |   deploy:
62 |     environment:
63 |       name: github-pages
64 |       url: ${{ steps.deployment.outputs.page_url }}
65 |     runs-on: ubuntu-latest
66 |     needs: build
67 |     steps:
68 |       - name: Deploy to GitHub Pages
69 |         id: deployment
70 |         uses: actions/deploy-pages@v4
71 | 


--------------------------------------------------------------------------------
/DOCKER_SETUP.md:
--------------------------------------------------------------------------------
 1 | # Docker Setup Guide
 2 | 
 3 | ## For End Users (Quick Start)
 4 | 
 5 | **Just use the pre-built images - no building required!**
 6 | 
 7 | ### Option 1: Docker Compose (Recommended)
 8 | 
 9 | ```bash
10 | # 1. Clone the repository
11 | git clone https://github.com/ls3-lab/QueryGym.git
12 | cd QueryGym
13 | 
14 | # 2. Set your API key
15 | export OPENAI_API_KEY="sk-..."
16 | 
17 | # 3. Run with Docker Compose
18 | docker compose run --rm querygym
19 | 
20 | # Or start Jupyter
21 | docker compose up jupyter
22 | ```
23 | 
24 | ### Option 2: Direct Docker Run
25 | 
26 | ```bash
27 | docker pull ghcr.io/ls3-lab/querygym:latest
28 | docker run -it --rm --gpus all ghcr.io/ls3-lab/querygym:latest
29 | ```
30 | 
31 | See the [full Docker guide](https://querygym.readthedocs.io/en/latest/user-guide/docker/) for more examples.
32 | 
33 | ---
34 | 
35 | ## For Developers (Building Locally)
36 | 
37 | If you need to build images locally for development:
38 | 
39 | ```bash
40 | # Build GPU image
41 | make build
42 | 
43 | # Build CPU image  
44 | make build-cpu
45 | 
46 | # Test builds
47 | make test
48 | 
49 | # Clean up
50 | make clean
51 | ```
52 | 
53 | The `Makefile` is for **local development only**. End users should use pre-built images.
54 | 
55 | ---
56 | 
57 | ## Files Overview
58 | 
59 | - **`compose.yml`** / **`docker-compose.yml`** - Uses pre-built GHCR images (for end users)
60 | - **`Dockerfile`** - GPU image definition (for building)
61 | - **`Dockerfile.cpu`** - CPU image definition (for building)
62 | - **`.dockerignore`** - Build optimization
63 | - **`Makefile`** - Development commands (for building locally)
64 | - **`examples/docker/`** - Ready-to-use examples with notebooks
65 | 
66 | ---
67 | 
68 | ## Images
69 | 
70 | - **GPU (default):** `ghcr.io/ls3-lab/querygym:latest`
71 | - **CPU (lightweight):** `ghcr.io/ls3-lab/querygym:cpu`
72 | 
73 | Images are automatically built and published on every release.
74 | 


--------------------------------------------------------------------------------
/querygym/methods/genqr.py:
--------------------------------------------------------------------------------
 1 | from __future__ import annotations
 2 | from ..core.base import BaseReformulator, QueryItem, ReformulationResult
 3 | from ..core.registry import register_method
 4 | from typing import List
 5 | 
 6 | @register_method("genqr")
 7 | class GENQR(BaseReformulator):
 8 |     """
 9 |     GenQR: Generate reformulations N times and concatenate.
10 |     
11 |     Pipeline:
12 |         1. Call LLM N times (default 5) to generate reformulations
13 |         2. Concatenate: query + reformulation1 + reformulation2 + ... + reformulationN
14 |     """
15 |     VERSION = "1.0"
16 | 
17 |     def reformulate(self, q: QueryItem, contexts=None) -> ReformulationResult:
18 |         # Get parameters
19 |         n_generations = int(self.cfg.params.get("n_generations", 5))
20 |         temperature = float(self.cfg.llm.get("temperature", 0.8))
21 |         max_tokens = int(self.cfg.llm.get("max_tokens", 256))
22 |         
23 |         # Generate reformulations N times (5 by default)
24 |         reformulations: List[str] = []
25 |         
26 |         for _ in range(n_generations):
27 |             msgs = self.prompts.render("genqr.v1", query=q.text)
28 |             out = self.llm.chat(msgs, temperature=temperature, max_tokens=max_tokens)
29 |             # Clean up the reformulation (remove quotes, extra whitespace)
30 |             reformulation = out.strip().strip('"').strip("'").replace("\n", " ")
31 |             reformulations.append(reformulation)
32 |         
33 |         # Concatenate: query + reformulation1 + reformulation2 + ... + reformulationN
34 |         reformulated = q.text + " " + " ".join(reformulations)
35 |         
36 |         return ReformulationResult(
37 |             q.qid, 
38 |             q.text, 
39 |             reformulated,
40 |             metadata={
41 |                 "n_generations": n_generations,
42 |                 "reformulations": reformulations
43 |             }
44 |         )
45 | 


--------------------------------------------------------------------------------
/querygym/config/defaults.yaml:
--------------------------------------------------------------------------------
 1 | llm:
 2 |   model: "gpt-4.1-mini"
 3 |   base_url: "${OPENAI_BASE_URL:-https://api.openai.com/v1}"
 4 |   api_key: "${OPENAI_API_KEY}"
 5 |   temperature: 1.0
 6 |   max_tokens: 1024
 7 | 
 8 | params:
 9 |   index: "msmarco-v1-passage"
10 |   retrieval_k: 10
11 |   k1: 0.9
12 |   b: 0.4
13 |   rm3: false
14 |   rocchio: false
15 |   rocchio_use_negative: false
16 |   disable_bm25_param: true
17 |   batch_size: 128
18 |   threads: 16
19 |   parallel: false  # Enable parallel generation for methods like MuGI
20 | 
21 | # Method-specific default parameters
22 | genqr:
23 |   llm:
24 |     temperature: 0.7
25 |     max_tokens: 256
26 |   params:
27 |     n_generations: 5  # Number of times to call LLM for keywords
28 | 
29 | query2e:
30 |   llm:
31 |     temperature: 0.7  
32 |     max_tokens: 256
33 |   params:
34 |     mode: "fs"  # "zs" (zero-shot) or "fs" (few-shot with 20 examples)
35 | 
36 | qa_expand:
37 |   llm:
38 |     temperature: 0.8
39 |     max_tokens: 256
40 |   params:
41 |     temperature_subq: 0.7  # Temperature for sub-question generation
42 |     temperature_answer: 0.7  # Temperature for answer generation
43 |     temperature_refine: 0.7  # Temperature for answer refinement
44 |     prompt_subq: "qa_expand.subq.v1"  # Prompt ID for sub-question generation
45 |     prompt_answer: "qa_expand.answer.v1"  # Prompt ID for answer generation
46 |     prompt_refine: "qa_expand.refine.v1"  # Prompt ID for answer refinement
47 | 
48 | genqr_ensemble:
49 |   llm:
50 |     temperature: 0.92  # Nucleus sampling (GenQREnsemble paper)
51 |     max_tokens: 256
52 |   params:
53 |     parallel: false  # Sequential generation by default (set true for parallel)
54 | 
55 | mugi:
56 |   llm:
57 |     temperature: 1.0  # High temperature for diversity (MuGI paper)
58 |     max_tokens: 1024  # Max tokens per pseudo-document generation
59 |   params:
60 |     num_docs: 5  # Number of pseudo-documents to generate (doc_gen in paper)
61 |     adaptive_times: 6  # Adaptive repetition divisor (default from paper)
62 |     parallel: false  # Sequential generation by default (set true for parallel)
63 |     mode: "zs"  # "zs" (zero-shot) or "fs" (few-shot)
64 | 
65 | seed: 17
66 | retries: 2


--------------------------------------------------------------------------------
/docs/about/license.md:
--------------------------------------------------------------------------------
 1 | # License
 2 | 
 3 | QueryGym is released under the **Apache License 2.0**.
 4 | 
 5 | ## Apache License 2.0
 6 | 
 7 | Copyright 2024-2025 QueryGym Contributors
 8 | 
 9 | Licensed under the Apache License, Version 2.0 (the "License");
10 | you may not use this file except in compliance with the License.
11 | You may obtain a copy of the License at
12 | 
13 |     http://www.apache.org/licenses/LICENSE-2.0
14 | 
15 | Unless required by applicable law or agreed to in writing, software
16 | distributed under the License is distributed on an "AS IS" BASIS,
17 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18 | See the License for the specific language governing permissions and
19 | limitations under the License.
20 | 
21 | See the full license text in the [LICENSE](https://github.com/ls3-lab/QueryGym/blob/main/LICENSE) file.
22 | 
23 | ## Third-Party Licenses
24 | 
25 | querygym depends on several open-source packages. Please refer to their respective licenses:
26 | 
27 | - **typer**: MIT License
28 | - **pyyaml**: MIT License
29 | - **openai**: Apache License 2.0
30 | - **tqdm**: MIT License / MPL-2.0
31 | 
32 | Optional dependencies:
33 | - **datasets** (HuggingFace): Apache License 2.0
34 | - **beir**: Apache License 2.0
35 | - **pyserini**: Apache License 2.0
36 | 
37 | ## Prompt Bank License
38 | 
39 | The prompts in `querygym/prompt_bank.yaml` are licensed under **CC-BY-4.0** (Creative Commons Attribution 4.0 International).
40 | 
41 | You are free to:
42 | - Share — copy and redistribute the prompts
43 | - Adapt — remix, transform, and build upon the prompts
44 | 
45 | Under the following terms:
46 | - Attribution — You must give appropriate credit to the original authors
47 | 
48 | ## Citation
49 | 
50 | If you use QueryGym in your research, please cite:
51 | 
52 | ```bibtex
53 | @misc{bigdeli2025querygymtoolkitreproduciblellmbased,
54 |       title={QueryGym: A Toolkit for Reproducible LLM-Based Query Reformulation}, 
55 |       author={Amin Bigdeli and Radin Hamidi Rad and Mert Incesu and Negar Arabzadeh and Charles L. A. Clarke and Ebrahim Bagheri},
56 |       year={2025},
57 |       eprint={2511.15996},
58 |       archivePrefix={arXiv},
59 |       primaryClass={cs.IR},
60 |       url={https://arxiv.org/abs/2511.15996}, 
61 | }
62 | ```
63 | 


--------------------------------------------------------------------------------
/docs/api/methods.md:
--------------------------------------------------------------------------------
  1 | # Methods API Reference
  2 | 
  3 | All methods inherit from `BaseReformulator` and provide the same interface.
  4 | 
  5 | ## Common Interface
  6 | 
  7 | All methods support:
  8 | 
  9 | ```python
 10 | # Single query reformulation
 11 | result = method.reformulate(query, contexts=None)
 12 | 
 13 | # Batch reformulation
 14 | results = method.reformulate_batch(queries, contexts=None)
 15 | ```
 16 | 
 17 | ## Available Methods
 18 | 
 19 | ### GenQR
 20 | 
 21 | Generic keyword expansion using LLM.
 22 | 
 23 | ```python
 24 | reformulator = qg.create_reformulator("genqr", model="gpt-4")
 25 | ```
 26 | 
 27 | **Requires Context:** No
 28 | 
 29 | ### GenQR Ensemble
 30 | 
 31 | Ensemble of multiple keyword expansion prompts.
 32 | 
 33 | ```python
 34 | reformulator = qg.create_reformulator(
 35 |     "genqr_ensemble",
 36 |     model="gpt-4",
 37 |     params={"repeat_query_weight": 3}
 38 | )
 39 | ```
 40 | 
 41 | **Requires Context:** No  
 42 | **Parameters:**
 43 | - `repeat_query_weight` (int): Number of query repetitions (default: 3)
 44 | 
 45 | ### Query2Doc
 46 | 
 47 | Generates pseudo-documents for the query.
 48 | 
 49 | ```python
 50 | reformulator = qg.create_reformulator("query2doc", model="gpt-4")
 51 | ```
 52 | 
 53 | **Requires Context:** No
 54 | 
 55 | ### QA Expand
 56 | 
 57 | Question-answer based expansion.
 58 | 
 59 | ```python
 60 | reformulator = qg.create_reformulator("qa_expand", model="gpt-4")
 61 | ```
 62 | 
 63 | **Requires Context:** No
 64 | 
 65 | ### MuGI
 66 | 
 67 | Multi-granularity information expansion.
 68 | 
 69 | ```python
 70 | reformulator = qg.create_reformulator("mugi", model="gpt-4")
 71 | ```
 72 | 
 73 | **Requires Context:** No
 74 | 
 75 | ### LameR
 76 | 
 77 | Context-based passage synthesis.
 78 | 
 79 | ```python
 80 | reformulator = qg.create_reformulator("lamer", model="gpt-4")
 81 | ```
 82 | 
 83 | **Requires Context:** Yes
 84 | 
 85 | ### Query2E
 86 | 
 87 | Query to entity expansion.
 88 | 
 89 | ```python
 90 | reformulator = qg.create_reformulator("query2e", model="gpt-4")
 91 | ```
 92 | 
 93 | **Requires Context:** No
 94 | 
 95 | ### CSQE
 96 | 
 97 | Context-based sentence extraction.
 98 | 
 99 | ```python
100 | reformulator = qg.create_reformulator("csqe", model="gpt-4")
101 | ```
102 | 
103 | **Requires Context:** Yes
104 | 


--------------------------------------------------------------------------------
/Dockerfile:
--------------------------------------------------------------------------------
 1 | # Default Dockerfile for querygym (GPU-enabled)
 2 | # This is the recommended version with GPU support
 3 | # For CPU-only (lightweight), use Dockerfile.cpu
 4 | # Base: NVIDIA CUDA official image for GPU support
 5 | 
 6 | # Use NVIDIA CUDA base image (Ubuntu-based)
 7 | FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
 8 | 
 9 | LABEL maintainer="querygym team <radin.h@gmail.com>"
10 | LABEL description="querygym: LLM-based Query Reformulation Toolkit with GPU support"
11 | LABEL org.opencontainers.image.source="https://github.com/ls3-lab/QueryGym"
12 | 
13 | # Install Python and system dependencies
14 | RUN apt-get update && apt-get install -y --no-install-recommends \
15 |     python3.10 \
16 |     python3.10-venv \
17 |     python3-pip \
18 |     openjdk-21-jre-headless \
19 |     build-essential \
20 |     git \
21 |     wget \
22 |     curl \
23 |     && rm -rf /var/lib/apt/lists/*
24 | 
25 | # Create symlinks for python
26 | RUN ln -sf /usr/bin/python3.10 /usr/bin/python && \
27 |     ln -sf /usr/bin/pip3 /usr/bin/pip
28 | 
29 | # Create virtual environment
30 | RUN python -m venv /opt/venv
31 | ENV PATH="/opt/venv/bin:$PATH"
32 | 
33 | # Upgrade pip
34 | RUN pip install --no-cache-dir --upgrade pip setuptools wheel
35 | 
36 | # Install PyTorch with CUDA support
37 | # CUDA 12.1 compatible wheels
38 | RUN pip install --no-cache-dir \
39 |     torch torchvision torchaudio \
40 |     --index-url https://download.pytorch.org/whl/cu121
41 | 
42 | # Install querygym with all dependencies
43 | RUN pip install --no-cache-dir \
44 |     querygym[all] \
45 |     pyserini \
46 |     jupyter \
47 |     ipython \
48 |     transformers \
49 |     sentence-transformers
50 | 
51 | # Set environment variables
52 | ENV PYTHONUNBUFFERED=1 \
53 |     PYTHONDONTWRITEBYTECODE=1 \
54 |     JAVA_HOME=/usr/lib/jvm/java-21-openjdk-amd64 \
55 |     CUDA_VISIBLE_DEVICES=0
56 | 
57 | # Create working directory
58 | WORKDIR /workspace
59 | 
60 | # Create directories
61 | RUN mkdir -p /workspace/data /workspace/outputs /workspace/notebooks
62 | 
63 | # Set up non-root user
64 | RUN useradd -m -u 1000 querygym && \
65 |     chown -R querygym:querygym /workspace
66 | 
67 | USER querygym
68 | 
69 | # Expose Jupyter port
70 | EXPOSE 8888
71 | 
72 | # Health check
73 | HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
74 |     CMD python -c "import querygym; import torch; assert torch.cuda.is_available()" || exit 1
75 | 
76 | CMD ["/bin/bash"]
77 | 


--------------------------------------------------------------------------------
/docs/api/adapters.md:
--------------------------------------------------------------------------------
 1 | # Adapters API Reference
 2 | 
 3 | ## Searcher Interface
 4 | 
 5 | The searcher interface provides a unified API for different search backends.
 6 | 
 7 | ### BaseSearcher
 8 | 
 9 | Abstract base class for all searcher implementations.
10 | 
11 | **Methods:**
12 | - `search(query: str, k: int = 10) -> List[SearchHit]` - Search for documents
13 | - `batch_search(queries: List[str], k: int = 10) -> Dict[str, List[SearchHit]]` - Batch search
14 | 
15 | ### SearchHit
16 | 
17 | Represents a single search result.
18 | 
19 | **Attributes:**
20 | - `docid: str` - Document ID
21 | - `score: float` - Relevance score
22 | - `text: Optional[str]` - Document text (if available)
23 | 
24 | ## Pyserini Adapter
25 | 
26 | ### PyseriniSearcher
27 | 
28 | Adapter for Pyserini's LuceneSearcher and LuceneImpactSearcher.
29 | 
30 | ```python
31 | from querygym.adapters import PyseriniSearcher
32 | 
33 | searcher = PyseriniSearcher(index_path="path/to/index")
34 | results = searcher.search("query text", k=10)
35 | ```
36 | 
37 | ## PyTerrier Adapter
38 | 
39 | ### PyTerrierSearcher
40 | 
41 | Adapter for PyTerrier retrievers.
42 | 
43 | ```python
44 | from querygym.adapters import PyTerrierSearcher
45 | import pyterrier as pt
46 | 
47 | pt.init()
48 | retriever = pt.BatchRetrieve(index, wmodel="BM25")
49 | searcher = PyTerrierSearcher(retriever)
50 | results = searcher.search("query text", k=10)
51 | ```
52 | 
53 | ## Wrapper Functions
54 | 
55 | ### wrap_pyserini_searcher
56 | 
57 | Wrap a Pyserini searcher instance.
58 | 
59 | ```python
60 | from querygym import wrap_pyserini_searcher
61 | from pyserini.search.lucene import LuceneSearcher
62 | 
63 | lucene_searcher = LuceneSearcher("path/to/index")
64 | searcher = wrap_pyserini_searcher(lucene_searcher)
65 | ```
66 | 
67 | ### wrap_pyterrier_retriever
68 | 
69 | Wrap a PyTerrier retriever instance.
70 | 
71 | ```python
72 | from querygym import wrap_pyterrier_retriever
73 | import pyterrier as pt
74 | 
75 | pt.init()
76 | retriever = pt.BatchRetrieve(index, wmodel="BM25")
77 | searcher = wrap_pyterrier_retriever(retriever)
78 | ```
79 | 
80 | ### wrap_custom_searcher
81 | 
82 | Wrap a custom search function.
83 | 
84 | ```python
85 | from querygym import wrap_custom_searcher
86 | 
87 | def my_search(query: str, k: int) -> List[dict]:
88 |     # Your custom search logic
89 |     return [{"docid": "doc1", "score": 0.9, "text": "..."}]
90 | 
91 | searcher = wrap_custom_searcher(my_search)
92 | ```
93 | 


--------------------------------------------------------------------------------
/docker-compose.yml:
--------------------------------------------------------------------------------
 1 | version: '3.8'
 2 | 
 3 | # QueryGym Docker Compose - Uses pre-built images from GitHub Container Registry
 4 | # No building required - just pull and run!
 5 | 
 6 | services:
 7 |   # GPU-enabled version (default)
 8 |   querygym:
 9 |     image: ghcr.io/ls3-lab/querygym:latest
10 |     container_name: querygym
11 |     deploy:
12 |       resources:
13 |         reservations:
14 |           devices:
15 |             - driver: nvidia
16 |               count: all
17 |               capabilities: [gpu]
18 |     environment:
19 |       - NVIDIA_VISIBLE_DEVICES=all
20 |       - OPENAI_API_KEY=${OPENAI_API_KEY}
21 |       - OPENAI_BASE_URL=${OPENAI_BASE_URL:-https://api.openai.com/v1}
22 |     volumes:
23 |       - ./data:/workspace/data
24 |       - ./outputs:/workspace/outputs
25 |       - ./notebooks:/workspace/notebooks
26 |     ports:
27 |       - "8888:8888"
28 |     stdin_open: true
29 |     tty: true
30 |     command: /bin/bash
31 | 
32 |   # CPU-only version (lightweight) - Uncomment if you don't have GPU
33 |   # querygym-cpu:
34 |   #   image: ghcr.io/ls3-lab/querygym:cpu
35 |   #   container_name: querygym-cpu
36 |   #   environment:
37 |   #     - OPENAI_API_KEY=${OPENAI_API_KEY}
38 |   #     - OPENAI_BASE_URL=${OPENAI_BASE_URL:-https://api.openai.com/v1}
39 |   #   volumes:
40 |   #     - ./data:/workspace/data
41 |   #     - ./outputs:/workspace/outputs
42 |   #     - ./notebooks:/workspace/notebooks
43 |   #   ports:
44 |   #     - "8889:8888"
45 |   #   stdin_open: true
46 |   #   tty: true
47 |   #   command: /bin/bash
48 | 
49 |   # Jupyter notebook server (GPU-enabled)
50 |   jupyter:
51 |     image: ghcr.io/ls3-lab/querygym:latest
52 |     container_name: querygym-jupyter
53 |     deploy:
54 |       resources:
55 |         reservations:
56 |           devices:
57 |             - driver: nvidia
58 |               count: all
59 |               capabilities: [gpu]
60 |     environment:
61 |       - NVIDIA_VISIBLE_DEVICES=all
62 |       - OPENAI_API_KEY=${OPENAI_API_KEY}
63 |       - OPENAI_BASE_URL=${OPENAI_BASE_URL:-https://api.openai.com/v1}
64 |     volumes:
65 |       - ./data:/workspace/data
66 |       - ./outputs:/workspace/outputs
67 |       - ./notebooks:/workspace/notebooks
68 |     ports:
69 |       - "8888:8888"
70 |     command: >
71 |       jupyter notebook
72 |       --ip=0.0.0.0
73 |       --port=8888
74 |       --no-browser
75 |       --allow-root
76 |       --NotebookApp.token=''
77 |       --NotebookApp.password=''
78 | 


--------------------------------------------------------------------------------
/querygym/core/llm.py:
--------------------------------------------------------------------------------
 1 | from __future__ import annotations
 2 | import os
 3 | from openai import OpenAI
 4 | from typing import List, Dict, Any
 5 | 
 6 | class OpenAICompatibleClient:
 7 |     """Client for OpenAI-compatible chat completion APIs.
 8 |     
 9 |     Supports any API that implements the OpenAI chat completions format,
10 |     including OpenAI, Azure OpenAI, local models via vLLM, etc.
11 |     
12 |     Attributes:
13 |         client: OpenAI client instance
14 |         model: Model identifier to use for completions
15 |         
16 |     Example:
17 |         >>> client = OpenAICompatibleClient(
18 |         ...     model="gpt-4",
19 |         ...     base_url="https://api.openai.com/v1",
20 |         ...     api_key="sk-..."
21 |         ... )
22 |         >>> response = client.chat([
23 |         ...     {"role": "system", "content": "You are helpful."},
24 |         ...     {"role": "user", "content": "Hello!"}
25 |         ... ])
26 |     """
27 |     
28 |     def __init__(self, model: str, base_url: str | None = None, api_key: str | None = None):
29 |         """Initialize the LLM client.
30 |         
31 |         Args:
32 |             model: Model identifier (e.g., "gpt-4", "gpt-3.5-turbo")
33 |             base_url: Optional base URL for the API. If None, uses OPENAI_BASE_URL
34 |                      environment variable or OpenAI's default.
35 |             api_key: Optional API key. If None, uses OPENAI_API_KEY environment variable.
36 |         """
37 |         self.client = OpenAI(base_url=base_url or os.getenv("OPENAI_BASE_URL", None),
38 |                              api_key=api_key or os.getenv("OPENAI_API_KEY"))
39 |         self.model = model
40 | 
41 |     def chat(self, messages: List[Dict[str, str]], **kw: Any) -> str:
42 |         """Send a chat completion request.
43 |         
44 |         Args:
45 |             messages: List of message dicts with 'role' and 'content' keys
46 |             **kw: Additional parameters (temperature, max_tokens, etc.)
47 |             
48 |         Returns:
49 |             Generated text response
50 |             
51 |         Example:
52 |             >>> response = client.chat(
53 |             ...     messages=[{"role": "user", "content": "Hello"}],
54 |             ...     temperature=0.8,
55 |             ...     max_tokens=256
56 |             ... )
57 |         """
58 |         resp = self.client.chat.completions.create(model=self.model, messages=messages, **kw)
59 |         return resp.choices[0].message.content or ""
60 | 


--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
 1 | .PHONY: help build build-cpu build-all test clean
 2 | 
 3 | # Makefile for QueryGym Docker Development
 4 | # This is for developers/contributors who need to build images locally
 5 | # 
 6 | # End users should use pre-built images:
 7 | #   docker pull ghcr.io/ls3-lab/querygym:latest
 8 | #   docker-compose up
 9 | 
10 | # Default target
11 | help:
12 | 	@echo "QueryGym Docker Development Commands"
13 | 	@echo "====================================="
14 | 	@echo ""
15 | 	@echo "⚠️  NOTE: This is for developers building images locally"
16 | 	@echo "   End users should use: docker-compose up"
17 | 	@echo ""
18 | 	@echo "Build Commands:"
19 | 	@echo "  make build          Build GPU image locally"
20 | 	@echo "  make build-cpu      Build CPU image locally"
21 | 	@echo "  make build-all      Build both images"
22 | 	@echo "  make test           Test the built images"
23 | 	@echo "  make clean          Remove built images"
24 | 	@echo ""
25 | 
26 | # Build targets
27 | build:
28 | 	@echo "Building GPU image locally..."
29 | 	@echo "Note: This may take 10-15 minutes"
30 | 	DOCKER_BUILDKIT=1 docker build -t querygym-dev:latest -f Dockerfile .
31 | 	@echo "✓ Built: querygym-dev:latest"
32 | 
33 | build-cpu:
34 | 	@echo "Building CPU image locally..."
35 | 	@echo "Note: This may take 5-10 minutes"
36 | 	DOCKER_BUILDKIT=1 docker build -t querygym-dev:cpu -f Dockerfile.cpu .
37 | 	@echo "✓ Built: querygym-dev:cpu"
38 | 
39 | build-all: build build-cpu
40 | 	@echo "✓ All images built successfully"
41 | 
42 | # Test the built images
43 | test:
44 | 	@echo "Testing locally built images..."
45 | 	@echo ""
46 | 	@echo "Testing querygym-dev:latest (GPU)..."
47 | 	docker run --rm querygym-dev:latest python -c "import querygym; print(f'✓ querygym {querygym.__version__}')"
48 | 	docker run --rm querygym-dev:latest python -c "import pyserini; print('✓ PySerini OK')"
49 | 	docker run --rm querygym-dev:latest python -c "import torch; print(f'✓ PyTorch {torch.__version__}')"
50 | 	@echo ""
51 | 	@echo "Testing querygym-dev:cpu..."
52 | 	docker run --rm querygym-dev:cpu python -c "import querygym; print(f'✓ querygym {querygym.__version__}')"
53 | 	docker run --rm querygym-dev:cpu python -c "import pyserini; print('✓ PySerini OK')"
54 | 	docker run --rm querygym-dev:cpu python -c "import torch; print(f'✓ PyTorch {torch.__version__}')"
55 | 	@echo ""
56 | 	@echo "✓ All tests passed!"
57 | 
58 | # Clean up locally built images
59 | clean:
60 | 	@echo "Removing locally built images..."
61 | 	docker rmi querygym-dev:latest querygym-dev:cpu 2>/dev/null || true
62 | 	@echo "✓ Cleanup complete!"
63 | 


--------------------------------------------------------------------------------
/Dockerfile.cpu:
--------------------------------------------------------------------------------
 1 | # CPU-only Dockerfile for querygym
 2 | # This is the alternate/lightweight version
 3 | # For GPU support (default), use the main Dockerfile or Dockerfile.gpu
 4 | # Base: python:3.10-slim (Debian-based) for best compatibility
 5 | 
 6 | # Stage 1: Builder stage
 7 | FROM python:3.10-slim AS builder
 8 | 
 9 | # Install build dependencies
10 | RUN apt-get update && apt-get install -y --no-install-recommends \
11 |     build-essential \
12 |     git \
13 |     wget \
14 |     && rm -rf /var/lib/apt/lists/*
15 | 
16 | # Create virtual environment
17 | RUN python -m venv /opt/venv
18 | ENV PATH="/opt/venv/bin:$PATH"
19 | 
20 | # Upgrade pip and install build tools
21 | RUN pip install --no-cache-dir --upgrade pip setuptools wheel
22 | 
23 | # Install PyTorch (CPU version for smaller image)
24 | # Using CPU-only wheel reduces image size by ~2GB vs CUDA version
25 | RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu
26 | 
27 | # Install querygym with all dependencies
28 | RUN pip install --no-cache-dir \
29 |     querygym[all] \
30 |     pyserini \
31 |     jupyter \
32 |     ipython
33 | 
34 | # Stage 2: Runtime stage
35 | FROM python:3.10-slim
36 | 
37 | LABEL maintainer="querygym team <radin.h@gmail.com>"
38 | LABEL description="querygym CPU: LLM-based Query Reformulation Toolkit (CPU-only)"
39 | LABEL org.opencontainers.image.source="https://github.com/ls3-lab/QueryGym"
40 | 
41 | # Install runtime dependencies (Java for PySerini)
42 | RUN apt-get update && apt-get install -y --no-install-recommends \
43 |     openjdk-21-jre-headless \
44 |     wget \
45 |     curl \
46 |     && rm -rf /var/lib/apt/lists/*
47 | 
48 | # Copy virtual environment from builder
49 | COPY --from=builder /opt/venv /opt/venv
50 | 
51 | # Set environment variables
52 | ENV PATH="/opt/venv/bin:$PATH" \
53 |     PYTHONUNBUFFERED=1 \
54 |     PYTHONDONTWRITEBYTECODE=1 \
55 |     JAVA_HOME=/usr/lib/jvm/java-21-openjdk-amd64
56 | 
57 | # Create working directory
58 | WORKDIR /workspace
59 | 
60 | # Create directories for data and outputs
61 | RUN mkdir -p /workspace/data /workspace/outputs /workspace/notebooks
62 | 
63 | # Set up a non-root user for security
64 | RUN useradd -m -u 1000 querygym && \
65 |     chown -R querygym:querygym /workspace
66 | 
67 | USER querygym
68 | 
69 | # Expose Jupyter port
70 | EXPOSE 8888
71 | 
72 | # Health check
73 | HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
74 |     CMD python -c "import querygym; import pyserini" || exit 1
75 | 
76 | # Default command: Start bash (users can override)
77 | CMD ["/bin/bash"]
78 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # Byte-compiled / optimized / DLL files
  2 | __pycache__/
  3 | *.py[cod]
  4 | *$py.class
  5 | 
  6 | # C extensions
  7 | *.so
  8 | 
  9 | # Distribution / packaging
 10 | .Python
 11 | build/
 12 | develop-eggs/
 13 | dist/
 14 | downloads/
 15 | eggs/
 16 | .eggs/
 17 | lib/
 18 | lib64/
 19 | parts/
 20 | sdist/
 21 | var/
 22 | wheels/
 23 | pip-wheel-metadata/
 24 | share/python-wheels/
 25 | *.egg-info/
 26 | .installed.cfg
 27 | *.egg
 28 | MANIFEST
 29 | 
 30 | # PyInstaller
 31 | *.manifest
 32 | *.spec
 33 | 
 34 | # Unit test / coverage reports
 35 | htmlcov/
 36 | .tox/
 37 | .nox/
 38 | .coverage
 39 | .coverage.*
 40 | .cache
 41 | nosetests.xml
 42 | coverage.xml
 43 | *.cover
 44 | *.py,cover
 45 | .hypothesis/
 46 | .pytest_cache/
 47 | queries_10.tsv
 48 | queries_all.tsv
 49 | output.tsv
 50 | 
 51 | # Translations
 52 | *.mo
 53 | *.pot
 54 | 
 55 | # Django stuff:
 56 | *.log
 57 | local_settings.py
 58 | db.sqlite3
 59 | db.sqlite3-journal
 60 | 
 61 | # Flask stuff:
 62 | instance/
 63 | .webassets-cache
 64 | 
 65 | # Scrapy stuff:
 66 | .scrapy
 67 | 
 68 | # Sphinx documentation
 69 | docs/_build/
 70 | 
 71 | # PyBuilder
 72 | target/
 73 | 
 74 | # Jupyter Notebook
 75 | .ipynb_checkpoints
 76 | 
 77 | # IPython
 78 | profile_default/
 79 | ipython_config.py
 80 | 
 81 | # pyenv
 82 | .python-version
 83 | 
 84 | # pipenv
 85 | Pipfile.lock
 86 | 
 87 | # PEP 582
 88 | __pypackages__/
 89 | 
 90 | # Celery stuff
 91 | celerybeat-schedule
 92 | celerybeat.pid
 93 | 
 94 | # SageMath parsed files
 95 | *.sage.py
 96 | 
 97 | # Environments
 98 | .env
 99 | .venv
100 | env/
101 | venv/
102 | ENV/
103 | env.bak/
104 | venv.bak/
105 | 
106 | # Spyder project settings
107 | .spyderproject
108 | .spyproject
109 | 
110 | # Rope project settings
111 | .ropeproject
112 | 
113 | # mkdocs documentation
114 | /site
115 | 
116 | # mypy
117 | .mypy_cache/
118 | .dmypy.json
119 | dmypy.json
120 | 
121 | # Pyre type checker
122 | .pyre/
123 | 
124 | # Ruff
125 | .ruff_cache/
126 | 
127 | # IDEs
128 | .vscode/
129 | .idea/
130 | *.swp
131 | *.swo
132 | *~
133 | .DS_Store
134 | 
135 | # Project specific
136 | reformulated_queries.tsv
137 | *.run
138 | *.eval.txt
139 | out/
140 | results/
141 | indices/
142 | datasets/
143 | scratch/
144 | query_expansion_results/
145 | dev_scripts/
146 | quick_independence_test.py
147 | test_adapters.py
148 | test_independence.py
149 | test_retriever_fix.py
150 | test_searcher_implementation.py
151 | test_standalone_final.py
152 | test_wrapper.py
153 | new_baselines.py
154 | outputs/
155 | docs/leaderboard-demo.html
156 | 
157 | test_examples.jsonl


--------------------------------------------------------------------------------
/docs/getting-started/installation.md:
--------------------------------------------------------------------------------
  1 | # Installation
  2 | 
  3 | ## Requirements
  4 | 
  5 | - Python 3.9 or higher
  6 | - pip package manager
  7 | 
  8 | ## Basic Installation
  9 | 
 10 | Install querygym from PyPI:
 11 | 
 12 | ```bash
 13 | pip install querygym
 14 | ```
 15 | 
 16 | This installs the core package with minimal dependencies:
 17 | - `typer` - CLI interface
 18 | - `pyyaml` - Prompt bank parsing
 19 | - `openai` - LLM client
 20 | - `tqdm` - Progress bars
 21 | 
 22 | ## Optional Dependencies
 23 | 
 24 | ### HuggingFace Datasets
 25 | 
 26 | For loading datasets from HuggingFace:
 27 | 
 28 | ```bash
 29 | pip install querygym[hf]
 30 | ```
 31 | 
 32 | ### BEIR Format Support
 33 | 
 34 | For working with BEIR benchmark datasets:
 35 | 
 36 | ```bash
 37 | pip install querygym[beir]
 38 | ```
 39 | 
 40 | ### Pyserini Integration
 41 | 
 42 | For using Pyserini search adapters:
 43 | 
 44 | ```bash
 45 | pip install querygym[pyserini]
 46 | ```
 47 | 
 48 | ### All Optional Features
 49 | 
 50 | Install everything:
 51 | 
 52 | ```bash
 53 | pip install querygym[all]
 54 | ```
 55 | 
 56 | ### Development Installation
 57 | 
 58 | For contributing to querygym:
 59 | 
 60 | ```bash
 61 | # Clone the repository
 62 | git clone https://github.com/ls3-lab/QueryGym.git
 63 | cd QueryGym
 64 | 
 65 | # Install in editable mode with dev dependencies
 66 | pip install -e .[dev]
 67 | ```
 68 | 
 69 | This includes:
 70 | - Testing tools (pytest, pytest-cov, pytest-mock)
 71 | - Code quality tools (black, ruff, mypy)
 72 | - Documentation tools (mkdocs, mkdocs-material, mkdocstrings)
 73 | 
 74 | ## Verify Installation
 75 | 
 76 | Check that querygym is installed correctly:
 77 | 
 78 | ```bash
 79 | # Check version
 80 | pip show querygym
 81 | 
 82 | # Test import
 83 | python -c "import querygym; print(querygym.__version__)"
 84 | 
 85 | # Test CLI
 86 | querygym --help
 87 | ```
 88 | 
 89 | ## Environment Setup
 90 | 
 91 | ### API Keys
 92 | 
 93 | querygym uses OpenAI-compatible APIs. Set your API key:
 94 | 
 95 | ```bash
 96 | export OPENAI_API_KEY="your-api-key-here"
 97 | ```
 98 | 
 99 | For custom endpoints:
100 | 
101 | ```bash
102 | export OPENAI_BASE_URL="https://your-custom-endpoint.com/v1"
103 | ```
104 | 
105 | ### Configuration File
106 | 
107 | Create a `.env` file in your project:
108 | 
109 | ```bash
110 | OPENAI_API_KEY=your-api-key-here
111 | OPENAI_BASE_URL=https://api.openai.com/v1  # Optional
112 | ```
113 | 
114 | ## Troubleshooting
115 | 
116 | ### Import Errors
117 | 
118 | If you get import errors for optional dependencies:
119 | 
120 | ```bash
121 | # Install the specific extra you need
122 | pip install querygym[beir]
123 | ```
124 | 
125 | ### Version Conflicts
126 | 
127 | If you have dependency conflicts:
128 | 
129 | ```bash
130 | # Create a fresh virtual environment
131 | python -m venv venv
132 | source venv/bin/activate  # On Windows: venv\Scripts\activate
133 | pip install querygym
134 | ```
135 | 
136 | ### Development Installation Issues
137 | 
138 | If editable install fails:
139 | 
140 | ```bash
141 | # Upgrade pip and setuptools
142 | pip install --upgrade pip setuptools wheel
143 | 
144 | # Try again
145 | pip install -e .[dev]
146 | ```
147 | 
148 | ## Next Steps
149 | 
150 | - [Quick Start Tutorial](quickstart.md)
151 | - [CLI Usage Guide](../user-guide/cli.md)
152 | - [API Reference](../api/core.md)
153 | 


--------------------------------------------------------------------------------
/mkdocs.yml:
--------------------------------------------------------------------------------
  1 | site_name: QueryGym Documentation
  2 | site_description: LLM-based Query Reformulation Toolkit
  3 | site_author: Amin Bigdeli, Radin Hamidi Rad, Mert Incesu, Negar Arabzadeh
  4 | site_url: https://querygym.readthedocs.io
  5 | 
  6 | repo_name: ls3-lab/QueryGym
  7 | repo_url: https://github.com/ls3-lab/QueryGym
  8 | edit_uri: edit/main/docs/
  9 | 
 10 | theme:
 11 |   name: material
 12 |   palette:
 13 |     # Light mode
 14 |     - scheme: default
 15 |       primary: indigo
 16 |       accent: indigo
 17 |       toggle:
 18 |         icon: material/brightness-7
 19 |         name: Switch to dark mode
 20 |     # Dark mode
 21 |     - scheme: slate
 22 |       primary: indigo
 23 |       accent: indigo
 24 |       toggle:
 25 |         icon: material/brightness-4
 26 |         name: Switch to light mode
 27 |   features:
 28 |     - navigation.tabs
 29 |     - navigation.sections
 30 |     - navigation.expand
 31 |     - navigation.top
 32 |     - search.suggest
 33 |     - search.highlight
 34 |     - content.code.copy
 35 |     - content.code.annotate
 36 | 
 37 | plugins:
 38 |   - search
 39 |   - mkdocstrings:
 40 |       handlers:
 41 |         python:
 42 |           paths: [.]
 43 |           options:
 44 |             docstring_style: google
 45 |             show_source: false
 46 |             show_root_heading: true
 47 |             show_root_full_path: false
 48 |             show_symbol_type_heading: true
 49 |             show_symbol_type_toc: true
 50 |             members_order: source
 51 |             show_signature_annotations: true
 52 | 
 53 | markdown_extensions:
 54 |   - pymdownx.highlight:
 55 |       anchor_linenums: true
 56 |   - pymdownx.inlinehilite
 57 |   - pymdownx.snippets
 58 |   - pymdownx.superfences
 59 |   - pymdownx.tabbed:
 60 |       alternate_style: true
 61 |   - admonition
 62 |   - pymdownx.details
 63 |   - pymdownx.emoji:
 64 |       emoji_index: !!python/name:material.extensions.emoji.twemoji
 65 |       emoji_generator: !!python/name:material.extensions.emoji.to_svg
 66 |   - attr_list
 67 |   - md_in_html
 68 |   - tables
 69 |   - toc:
 70 |       permalink: true
 71 | 
 72 | nav:
 73 |   - Home: index.md
 74 |   - Getting Started:
 75 |       - Installation: getting-started/installation.md
 76 |       - Quick Start: getting-started/quickstart.md
 77 |   - User Guide:
 78 |       - CLI Usage: user-guide/cli.md
 79 |       - Loading Datasets: user-guide/datasets.md
 80 |       - Query Reformulation: user-guide/reformulation.md
 81 |       - Methods Reference: user-guide/methods-reference.md
 82 |       - Docker Guide: user-guide/docker.md
 83 |       - Searcher Interface: user-guide/searcher.md
 84 |       - Prompt Bank: user-guide/prompts.md
 85 |   - API Reference:
 86 |       - Core: api/core.md
 87 |       - Methods: api/methods.md
 88 |       - Data Loaders: api/loaders.md
 89 |       - Adapters: api/adapters.md
 90 |   - Development:
 91 |       - Contributing: development/contributing.md
 92 |       - Development Guide: development/guide.md
 93 |   - About:
 94 |       - Changelog: about/changelog.md
 95 |       - License: about/license.md
 96 | 
 97 | extra:
 98 |   social:
 99 |     - icon: fontawesome/brands/github
100 |       link: https://github.com/ls3-lab/QueryGym
101 |     - icon: fontawesome/brands/python
102 |       link: https://pypi.org/project/querygym/
103 | 
104 | copyright: Copyright &copy; 2025 querygym Team
105 | 


--------------------------------------------------------------------------------
/pyproject.toml:
--------------------------------------------------------------------------------
  1 | [build-system]
  2 | requires = ["setuptools>=68", "wheel"]
  3 | build-backend = "setuptools.build_meta"
  4 | 
  5 | [project]
  6 | name = "querygym"
  7 | version = "0.1.5"
  8 | description = "LLM-based Query Reformulation Toolkit"
  9 | readme = "README.md"
 10 | requires-python = ">=3.9"
 11 | license = {text = "Apache-2.0"}
 12 | authors = [
 13 |   {name = "Radin Hamidi Rad", email = "radin.h@gmail.com"},
 14 |   {name = "Amin Bigdeli", email = "aminbigdeli97@gmail.com"},
 15 |   {name = "Mert Incesu", email = "mert.incesu03@gmail.com"},
 16 |   {name = "Negar Arabzadeh", email = "ngr.arabzadeh@gmail.com"}
 17 | ]
 18 | keywords = ["information-retrieval", "query-reformulation", "llm", "nlp"]
 19 | classifiers = [
 20 |     "Development Status :: 3 - Alpha",
 21 |     "Intended Audience :: Science/Research",
 22 |     "License :: OSI Approved :: Apache Software License",
 23 |     "Programming Language :: Python :: 3",
 24 |     "Programming Language :: Python :: 3.9",
 25 |     "Programming Language :: Python :: 3.10",
 26 |     "Programming Language :: Python :: 3.11",
 27 |     "Programming Language :: Python :: 3.12",
 28 |     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 29 | ]
 30 | 
 31 | dependencies = [
 32 |   "typer>=0.12.3",
 33 |   "pyyaml>=6.0.1",
 34 |   "openai>=1.40.0",
 35 |   "tqdm>=4.65.0",
 36 | ]
 37 | 
 38 | [project.optional-dependencies]
 39 | # Data loading extras
 40 | hf = ["datasets>=2.20.0"]
 41 | beir = ["beir>=2.0.0"]
 42 | pyserini = ["pyserini>=0.22.0"]
 43 | 
 44 | # Development dependencies
 45 | dev = [
 46 |   # Testing
 47 |   "pytest>=7.4.0",
 48 |   "pytest-cov>=4.1.0",
 49 |   "pytest-mock>=3.11.0",
 50 |   
 51 |   # Code quality
 52 |   "black>=23.0.0",
 53 |   "ruff>=0.1.0",
 54 |   "mypy>=1.5.0",
 55 |   
 56 |   # Documentation
 57 |   "mkdocs>=1.5.0",
 58 |   "mkdocs-material>=9.0.0",
 59 |   "mkdocstrings[python]>=0.23.0",
 60 | ]
 61 | 
 62 | # All extras combined
 63 | all = [
 64 |   "datasets>=2.20.0",
 65 |   "beir>=2.0.0",
 66 |   "pyserini>=0.22.0",
 67 | ]
 68 | 
 69 | [project.scripts]
 70 | querygym = "querygym.cli:app"
 71 | 
 72 | [project.urls]
 73 | Homepage = "https://github.com/ls3-lab/QueryGym"
 74 | Documentation = "https://github.com/ls3-lab/QueryGym/blob/main/README.md"
 75 | Repository = "https://github.com/ls3-lab/QueryGym"
 76 | Issues = "https://github.com/ls3-lab/QueryGym/issues"
 77 | 
 78 | [tool.setuptools.packages.find]
 79 | where = ["."]
 80 | include = ["querygym*"]
 81 | 
 82 | [tool.setuptools.package-data]
 83 | querygym = ["prompt_bank.yaml", "config/*.yaml"]
 84 | 
 85 | # Black configuration
 86 | [tool.black]
 87 | line-length = 100
 88 | target-version = ["py39", "py310", "py311"]
 89 | include = '\.pyi?$'
 90 | 
 91 | # Ruff configuration
 92 | [tool.ruff]
 93 | line-length = 100
 94 | target-version = "py39"
 95 | select = ["E", "F", "I", "N", "W"]
 96 | ignore = ["E501"]  # Line too long (handled by black)
 97 | 
 98 | # Pytest configuration
 99 | [tool.pytest.ini_options]
100 | testpaths = ["tests"]
101 | python_files = ["test_*.py"]
102 | python_classes = ["Test*"]
103 | python_functions = ["test_*"]
104 | addopts = [
105 |     "--verbose",
106 |     "--cov=querygym",
107 |     "--cov-report=term-missing",
108 |     "--cov-report=html",
109 | ]
110 | 
111 | # MyPy configuration
112 | [tool.mypy]
113 | python_version = "3.9"
114 | warn_return_any = true
115 | warn_unused_configs = true
116 | disallow_untyped_defs = false
117 | ignore_missing_imports = true
118 | 


--------------------------------------------------------------------------------
/.github/workflows/docker.yml:
--------------------------------------------------------------------------------
 1 | name: Build and Push Docker Images
 2 | 
 3 | on:
 4 |   release:
 5 |     types: [published]
 6 |   push:
 7 |     branches:
 8 |       - main
 9 |     paths:
10 |       - 'Dockerfile*'
11 |       - '.github/workflows/docker.yml'
12 |   workflow_dispatch:
13 | 
14 | env:
15 |   REGISTRY: ghcr.io
16 |   IMAGE_NAME: ${{ github.repository }}
17 | 
18 | jobs:
19 |   build-and-push:
20 |     runs-on: ubuntu-latest
21 |     permissions:
22 |       contents: read
23 |       packages: write
24 | 
25 |     strategy:
26 |       matrix:
27 |         include:
28 |           - dockerfile: Dockerfile
29 |             tag-suffix: ""
30 |             extra-tags: ""
31 |             platforms: linux/amd64
32 |             description: "GPU-enabled (default)"
33 |           - dockerfile: Dockerfile.cpu
34 |             tag-suffix: "-cpu"
35 |             extra-tags: "type=raw,value=cpu,enable={{is_default_branch}}"
36 |             platforms: linux/amd64,linux/arm64
37 |             description: "CPU-only (lightweight)"
38 | 
39 |     steps:
40 |       - name: Checkout repository
41 |         uses: actions/checkout@v4
42 | 
43 |       - name: Set up QEMU
44 |         uses: docker/setup-qemu-action@v3
45 | 
46 |       - name: Set up Docker Buildx
47 |         uses: docker/setup-buildx-action@v3
48 | 
49 |       - name: Log in to GitHub Container Registry
50 |         uses: docker/login-action@v3
51 |         with:
52 |           registry: ${{ env.REGISTRY }}
53 |           username: ${{ github.actor }}
54 |           password: ${{ secrets.GITHUB_TOKEN }}
55 | 
56 |       - name: Extract metadata
57 |         id: meta
58 |         uses: docker/metadata-action@v5
59 |         with:
60 |           images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
61 |           tags: |
62 |             type=semver,pattern={{version}}${{ matrix.tag-suffix }}
63 |             type=semver,pattern={{major}}.{{minor}}${{ matrix.tag-suffix }}
64 |             type=sha,prefix=sha-,suffix=${{ matrix.tag-suffix }},enable=${{ github.event_name != 'release' }}
65 |             type=raw,value=latest${{ matrix.tag-suffix }},enable={{is_default_branch}}
66 |             ${{ matrix.extra-tags }}
67 |           labels: |
68 |             org.opencontainers.image.title=QueryGym
69 |             org.opencontainers.image.description=LLM-based Query Reformulation Toolkit
70 |             org.opencontainers.image.vendor=QueryGym Team
71 |             org.opencontainers.image.licenses=Apache-2.0
72 | 
73 |       - name: Build and push Docker image
74 |         uses: docker/build-push-action@v5
75 |         with:
76 |           context: .
77 |           file: ${{ matrix.dockerfile }}
78 |           platforms: ${{ matrix.platforms }}
79 |           push: true
80 |           tags: ${{ steps.meta.outputs.tags }}
81 |           labels: ${{ steps.meta.outputs.labels }}
82 |           cache-from: type=gha
83 |           cache-to: type=gha,mode=max
84 |           
85 |       - name: Generate image summary
86 |         run: |
87 |           echo "### 🐳 Docker Image Published" >> $GITHUB_STEP_SUMMARY
88 |           echo "" >> $GITHUB_STEP_SUMMARY
89 |           echo "**Image:** \`${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}${{ matrix.tag-suffix }}\`" >> $GITHUB_STEP_SUMMARY
90 |           echo "**Description:** ${{ matrix.description }}" >> $GITHUB_STEP_SUMMARY
91 |           echo "" >> $GITHUB_STEP_SUMMARY
92 |           echo "Pull command:" >> $GITHUB_STEP_SUMMARY
93 |           echo "\`\`\`bash" >> $GITHUB_STEP_SUMMARY
94 |           echo "docker pull ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest${{ matrix.tag-suffix }}" >> $GITHUB_STEP_SUMMARY
95 |           echo "\`\`\`" >> $GITHUB_STEP_SUMMARY
96 | 


--------------------------------------------------------------------------------
/examples/README.md:
--------------------------------------------------------------------------------
  1 | # QueryGym Examples
  2 | 
  3 | This directory contains practical examples and tutorials for using QueryGym in different scenarios.
  4 | 
  5 | ## 📁 Directory Structure
  6 | 
  7 | ### 1. [snippets](snippets/)
  8 | **Quick code snippets** demonstrating core QueryGym functionality.
  9 | 
 10 | Perfect for:
 11 | - Learning the basics
 12 | - Quick reference
 13 | - Copy-paste solutions
 14 | 
 15 | Examples include:
 16 | - Loading data from files
 17 | - Using format loaders (BEIR, MS MARCO)
 18 | - Query reformulation
 19 | - Context-based reformulation
 20 | - Using search adapters
 21 | 
 22 | ### 2. [docker](docker/)
 23 | **Docker-based examples** for containerized workflows.
 24 | 
 25 | Perfect for:
 26 | - Quick start without installation
 27 | - Reproducible environments
 28 | - Jupyter notebook development
 29 | 
 30 | Includes:
 31 | - Ready-to-use `docker-compose.yml` (Compose V2 compatible)
 32 | - Interactive Jupyter notebook
 33 | - Quick setup guide
 34 | 
 35 | ### 3. [querygym_pyserini](querygym_pyserini/)
 36 | **End-to-end retrieval workflows** integrating QueryGym with Pyserini.
 37 | 
 38 | Perfect for:
 39 | - Complete IR experiments
 40 | - Indexing and retrieval pipelines
 41 | - Evaluation workflows
 42 | 
 43 | Includes:
 44 | - Building indexes
 45 | - Query reformulation + retrieval
 46 | - Evaluation with trec_eval
 47 | - Complete experiment scripts
 48 | 
 49 | ## 🚀 Quick Start
 50 | 
 51 | ### Option 1: Run with Docker (Easiest)
 52 | ```bash
 53 | cd docker/
 54 | docker compose run --rm querygym
 55 | ```
 56 | 
 57 | ### Option 2: Run Python Snippets
 58 | ```bash
 59 | # Install QueryGym first
 60 | pip install querygym
 61 | 
 62 | # Run any snippet
 63 | cd snippets/
 64 | python 01_load_from_file.py
 65 | ```
 66 | 
 67 | ### Option 3: Full Retrieval Pipeline
 68 | ```bash
 69 | cd querygym_pyserini/
 70 | # Follow the README for complete setup
 71 | ```
 72 | 
 73 | ## 📚 Learning Path
 74 | 
 75 | **Beginner:**
 76 | - Start with [snippets](snippets/) to learn the basics
 77 | - Try [docker](docker/) for a containerized environment
 78 | 
 79 | **Intermediate:**
 80 | - Explore different reformulation methods in snippets
 81 | - Run Jupyter notebooks in Docker
 82 | 
 83 | **Advanced:**
 84 | - Build complete pipelines with [querygym_pyserini](querygym_pyserini/)
 85 | - Integrate with your own retrieval systems
 86 | 
 87 | ## 🔗 Additional Resources
 88 | 
 89 | - **Documentation:** [https://querygym.readthedocs.io](https://querygym.readthedocs.io/)
 90 | - **Docker Setup:** [DOCKER_SETUP.md](../../DOCKER_SETUP.md) in the root directory
 91 | - **API Reference:** [https://querygym.readthedocs.io/en/latest/api/core/](https://querygym.readthedocs.io/en/latest/api/core/)
 92 | - **GitHub:** [https://github.com/ls3-lab/QueryGym](https://github.com/ls3-lab/QueryGym)
 93 | 
 94 | ## 💡 Tips
 95 | 
 96 | - **Start simple:** Begin with basic snippets before complex workflows
 97 | - **Use Docker:** For quick experimentation without setup hassles
 98 | - **Check READMEs:** Each subdirectory has its own detailed README
 99 | - **Modify examples:** All examples are meant to be adapted to your needs
100 | 
101 | ## 🤝 Contributing Examples
102 | 
103 | Have a useful example? We'd love to include it! Please:
104 | 1. Add it to the appropriate subdirectory
105 | 2. Include clear comments and documentation
106 | 3. Test it works with the latest QueryGym version
107 | 4. Submit a pull request
108 | 
109 | See the [Contributing section](../README.md#contributing) in the main README for general contribution guidelines.
110 | 
111 | ## ❓ Need Help?
112 | 
113 | - Check the [documentation](https://querygym.readthedocs.io/)
114 | - Open an [issue](https://github.com/ls3-lab/QueryGym/issues)
115 | - Join discussions on GitHub
116 | 


--------------------------------------------------------------------------------
/docs/index.md:
--------------------------------------------------------------------------------
  1 | <div align="center">
  2 |   <img src="querygym-logo.png" alt="QueryGym Logo" width="600">
  3 | </div>
  4 | 
  5 | # QueryGym
  6 | 
  7 | A lightweight, reproducible toolkit for **LLM-based query reformulation**.
  8 | 
  9 | 
 10 | [![Publish to PyPI](https://github.com/ls3-lab/QueryGym/actions/workflows/publish.yml/badge.svg)](https://github.com/ls3-lab/QueryGym/actions/workflows/publish.yml)
 11 | [![Build and Push Docker Images](https://github.com/ls3-lab/QueryGym/actions/workflows/docker.yml/badge.svg)](https://github.com/ls3-lab/QueryGym/actions/workflows/docker.yml)
 12 | [![PyPI version](https://badge.fury.io/py/querygym.svg)](https://pypi.org/project/querygym/)
 13 | ![PyPI - Downloads](https://img.shields.io/pypi/dw/querygym?color=blueviolet&label=downloads)
 14 | [![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
 15 | [![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 16 | 
 17 | ## Features
 18 | 
 19 | - **Single Prompt Bank** (YAML) with metadata
 20 | - **Simple DataLoader**: Dependency-free file loading for queries, qrels, and contexts
 21 | - **Format Loaders**: Optional BEIR and MS MARCO format loaders
 22 | - **OpenAI-compatible** LLM client (works with any OpenAI API–compatible endpoint)
 23 | - **Pyserini** optional: either pass contexts (JSONL) or pass a retriever instance to build contexts
 24 | - Export-only: emits reformulated queries; optionally generates a **bash** script for Pyserini + `trec_eval`
 25 | 
 26 | ## Quick Example
 27 | 
 28 | ```python
 29 | import querygym as qg
 30 | 
 31 | # Load data
 32 | queries = qg.load_queries("queries.tsv")
 33 | qrels = qg.load_qrels("qrels.txt")
 34 | 
 35 | # Create reformulator
 36 | reformulator = qg.create_reformulator("genqr_ensemble", model="gpt-4")
 37 | 
 38 | # Reformulate
 39 | results = reformulator.reformulate_batch(queries)
 40 | 
 41 | # Save
 42 | qg.DataLoader.save_queries(
 43 |     [qg.QueryItem(r.qid, r.reformulated) for r in results],
 44 |     "reformulated.tsv"
 45 | )
 46 | ```
 47 | 
 48 | ## Installation
 49 | 
 50 | ### Install from PyPI
 51 | ```bash
 52 | pip install querygym
 53 | ```
 54 | 
 55 | ### Use Docker (Quick Start)
 56 | ```bash
 57 | # Pull pre-built image
 58 | docker pull ghcr.io/ls3-lab/querygym:latest
 59 | 
 60 | # Run with Docker Compose
 61 | docker compose run --rm querygym
 62 | ```
 63 | 
 64 | See the [Docker Guide](user-guide/docker.md) for detailed setup and usage.
 65 | 
 66 | For optional features:
 67 | 
 68 | ```bash
 69 | # With HuggingFace datasets support
 70 | pip install querygym[hf]
 71 | 
 72 | # With BEIR format support
 73 | pip install querygym[beir]
 74 | 
 75 | # With Pyserini adapter
 76 | pip install querygym[pyserini]
 77 | 
 78 | # All optional features
 79 | pip install querygym[all]
 80 | 
 81 | # Development dependencies
 82 | pip install querygym[dev]
 83 | ```
 84 | 
 85 | ## Documentation
 86 | 
 87 | - [Installation Guide](getting-started/installation.md)
 88 | - [Quick Start Tutorial](getting-started/quickstart.md)
 89 | - [CLI Usage](user-guide/cli.md)
 90 | - [API Reference](api/core.md)
 91 | 
 92 | > 📊 **Looking for benchmarks?** Visit the [Leaderboard](https://ls3-lab.github.io/QueryGym/leaderboard.html) on our main site.
 93 | 
 94 | ## Citation
 95 | 
 96 | If you use QueryGym in your research, please cite:
 97 | 
 98 | ```bibtex
 99 | @misc{bigdeli2025querygymtoolkitreproduciblellmbased,
100 |       title={QueryGym: A Toolkit for Reproducible LLM-Based Query Reformulation}, 
101 |       author={Amin Bigdeli and Radin Hamidi Rad and Mert Incesu and Negar Arabzadeh and Charles L. A. Clarke and Ebrahim Bagheri},
102 |       year={2025},
103 |       eprint={2511.15996},
104 |       archivePrefix={arXiv},
105 |       primaryClass={cs.IR},
106 |       url={https://arxiv.org/abs/2511.15996}, 
107 | }
108 | ```
109 | 
110 | ## License
111 | 
112 | Apache License 2.0 - see [LICENSE](about/license.md) for details.
113 | 


--------------------------------------------------------------------------------
/docs/user-guide/reformulation.md:
--------------------------------------------------------------------------------
  1 | # Query Reformulation Methods
  2 | 
  3 | querygym provides several state-of-the-art query reformulation methods.
  4 | 
  5 | ## Available Methods
  6 | 
  7 | ### GenQR (Generic Query Reformulation)
  8 | 
  9 | Simple keyword expansion using LLM.
 10 | 
 11 | ```python
 12 | import querygym as qg
 13 | 
 14 | reformulator = qg.create_reformulator("genqr", model="gpt-4")
 15 | result = reformulator.reformulate(qg.QueryItem("q1", "neural networks"))
 16 | ```
 17 | 
 18 | ### GenQR Ensemble
 19 | 
 20 | Ensemble of multiple keyword expansion prompts for better coverage.
 21 | 
 22 | ```python
 23 | reformulator = qg.create_reformulator(
 24 |     "genqr_ensemble",
 25 |     model="gpt-4",
 26 |     params={"repeat_query_weight": 3}
 27 | )
 28 | ```
 29 | 
 30 | **Parameters:**
 31 | - `repeat_query_weight` (int): Number of times to repeat original query (default: 3)
 32 | 
 33 | ### Query2Doc
 34 | 
 35 | Generates pseudo-documents relevant to the query.
 36 | 
 37 | ```python
 38 | reformulator = qg.create_reformulator("query2doc", model="gpt-4")
 39 | ```
 40 | 
 41 | Supports both zero-shot and chain-of-thought variants.
 42 | 
 43 | ### QA Expand
 44 | 
 45 | Decomposes query into sub-questions, generates answers, and refines.
 46 | 
 47 | ```python
 48 | reformulator = qg.create_reformulator("qa_expand", model="gpt-4")
 49 | ```
 50 | 
 51 | ### MuGI
 52 | 
 53 | Multi-granularity information expansion.
 54 | 
 55 | ```python
 56 | reformulator = qg.create_reformulator("mugi", model="gpt-4")
 57 | ```
 58 | 
 59 | ### LameR
 60 | 
 61 | Context-based passage synthesis using retrieved documents.
 62 | 
 63 | ```python
 64 | # Load contexts
 65 | contexts = qg.load_contexts("contexts.jsonl")
 66 | 
 67 | # Create reformulator
 68 | reformulator = qg.create_reformulator("lamer", model="gpt-4")
 69 | 
 70 | # Reformulate with contexts
 71 | results = reformulator.reformulate_batch(queries, contexts=contexts)
 72 | ```
 73 | 
 74 | **Note:** LameR requires contexts from initial retrieval.
 75 | 
 76 | ### Query2E
 77 | 
 78 | Query to entity expansion.
 79 | 
 80 | ```python
 81 | reformulator = qg.create_reformulator("query2e", model="gpt-4")
 82 | ```
 83 | 
 84 | ### CSQE
 85 | 
 86 | Context-based sentence extraction from retrieved documents.
 87 | 
 88 | ```python
 89 | # Requires contexts
 90 | contexts = qg.load_contexts("contexts.jsonl")
 91 | 
 92 | reformulator = qg.create_reformulator("csqe", model="gpt-4")
 93 | results = reformulator.reformulate_batch(queries, contexts=contexts)
 94 | ```
 95 | 
 96 | ## Method Comparison
 97 | 
 98 | | Method | Requires Context | Type | Best For |
 99 | |--------|-----------------|------|----------|
100 | | genqr | No | Keyword expansion | General queries |
101 | | genqr_ensemble | No | Keyword expansion | Robust expansion |
102 | | query2doc | No | Pseudo-document | Dense retrieval |
103 | | qa_expand | No | QA-based | Complex queries |
104 | | mugi | No | Multi-granular | Diverse expansion |
105 | | lamer | Yes | Context synthesis | Re-ranking |
106 | | query2e | No | Entity expansion | Entity queries |
107 | | csqe | Yes | Sentence extraction | Precision-focused |
108 | 
109 | ## Custom Parameters
110 | 
111 | All methods support custom parameters:
112 | 
113 | ```python
114 | reformulator = qg.create_reformulator(
115 |     "genqr_ensemble",
116 |     model="gpt-4",
117 |     params={
118 |         "repeat_query_weight": 5,
119 |         "temperature": 0.7
120 |     },
121 |     llm_config={
122 |         "temperature": 0.8,
123 |         "max_tokens": 512
124 |     }
125 | )
126 | ```
127 | 
128 | ## Batch Processing
129 | 
130 | Process multiple queries efficiently:
131 | 
132 | ```python
133 | queries = qg.load_queries("queries.tsv")
134 | reformulator = qg.create_reformulator("genqr", model="gpt-4")
135 | 
136 | # Batch reformulation with progress bar
137 | results = reformulator.reformulate_batch(queries)
138 | ```
139 | 
140 | ## Using Custom Prompts
141 | 
142 | See [Prompt Bank](prompts.md) for details on customizing prompts.
143 | 
144 | ## Next Steps
145 | 
146 | - [Prompt Bank Documentation](prompts.md)
147 | - [API Reference](../api/methods.md)
148 | 


--------------------------------------------------------------------------------
/docs/about/changelog.md:
--------------------------------------------------------------------------------
  1 | # Changelog
  2 | 
  3 | All notable changes to querygym will be documented in this file.
  4 | 
  5 | The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
  6 | and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
  7 | 
  8 | ## [0.1.4] - 2025-11-21
  9 | 
 10 | ### Added
 11 | - Docker support with pre-built images on GitHub Container Registry (GHCR)
 12 |   - GPU-enabled image (`ghcr.io/ls3-lab/querygym:latest`)
 13 |   - CPU-only image (`ghcr.io/ls3-lab/querygym:cpu`)
 14 |   - Multi-platform support (linux/amd64, linux/arm64 for CPU)
 15 | - Docker Compose configuration for easy setup
 16 | - Comprehensive Docker documentation and examples
 17 | - GitHub Actions workflow for automated Docker image building and publishing
 18 | - Example notebooks and scripts in `examples/` directory
 19 |   - Code snippets for quick reference
 20 |   - Docker examples with Jupyter notebooks
 21 |   - QueryGym + Pyserini integration examples
 22 | - DOCKER_SETUP.md quick reference guide
 23 | - Makefile for local Docker development
 24 | 
 25 | ### Changed
 26 | - **License changed from MIT to Apache License 2.0**
 27 | - Reorganized examples into three categories: snippets, docker, querygym_pyserini
 28 | - Updated to OpenJDK 21 (from 17) in Docker images
 29 | - Migrated to Docker Compose V2 syntax (`docker compose` instead of `docker-compose`)
 30 | - Improved Docker image tagging strategy (separate tags for GPU and CPU)
 31 | - Enhanced README with Docker installation options and citation information
 32 | - Updated all documentation to use modern Docker Compose commands
 33 | 
 34 | ### Fixed
 35 | - Docker build issues with OpenJDK availability in Debian Trixie
 36 | - Docker Compose tag conflicts between GPU and CPU images
 37 | - Inconsistent command examples across documentation
 38 | 
 39 | ## [0.1.3] - 2025-11-19
 40 | 
 41 | ### Added
 42 | - Apache License 2.0
 43 | - Citation information with arXiv paper reference in README
 44 | 
 45 | ### Changed
 46 | - Package metadata updated with Apache 2.0 license
 47 | 
 48 | ## [0.1.2] - 2025-11-18
 49 | 
 50 | ### Changed
 51 | - Updated author information with email addresses for Mert Incesu and Negar Arabzadeh
 52 | - Version bump for metadata updates
 53 | 
 54 | ## [0.1.1] - 2025-11-18
 55 | 
 56 | ### Added
 57 | - Added `tqdm` as core dependency for progress bars
 58 | - Complete Read the Docs documentation setup
 59 | - MkDocs configuration with Material theme
 60 | - Comprehensive API reference documentation
 61 | - GitHub Actions workflow for automated PyPI publishing
 62 | 
 63 | ### Changed
 64 | - Updated package name from `queryGym` to `querygym` (lowercase)
 65 | - Updated author information in package metadata
 66 | - Improved dependency management
 67 | 
 68 | ### Fixed
 69 | - Fixed missing `tqdm` dependency in core package
 70 | 
 71 | ## [0.1.0] - 2025-11-17
 72 | 
 73 | ### Added
 74 | - Initial release of querygym
 75 | - Core reformulation framework
 76 | - Eight query reformulation methods:
 77 |   - GenQR (generic keyword expansion)
 78 |   - GenQR Ensemble
 79 |   - Query2Doc
 80 |   - QA Expand
 81 |   - MuGI
 82 |   - LameR
 83 |   - Query2E
 84 |   - CSQE
 85 | - Prompt bank system with YAML configuration
 86 | - DataLoader for queries, qrels, and contexts
 87 | - BEIR and MS MARCO format loaders
 88 | - Pyserini and PyTerrier adapters
 89 | - CLI interface with `querygym` command
 90 | - OpenAI-compatible LLM client
 91 | - Comprehensive test suite
 92 | - Example scripts and documentation
 93 | 
 94 | ### Features
 95 | - Dependency-free core data loading
 96 | - Optional integrations (BEIR, Pyserini, HuggingFace)
 97 | - Flexible LLM client supporting any OpenAI-compatible API
 98 | - Batch processing with progress bars
 99 | - Context-based reformulation support
100 | - Export to TSV/JSONL formats
101 | - Script generation for Pyserini + trec_eval
102 | 
103 | [Unreleased]: https://github.com/ls3-lab/QueryGym/compare/v0.1.3...HEAD
104 | [0.1.3]: https://github.com/ls3-lab/QueryGym/compare/v0.1.2...v0.1.3
105 | [0.1.2]: https://github.com/ls3-lab/QueryGym/compare/v0.1.1...v0.1.2
106 | [0.1.1]: https://github.com/ls3-lab/QueryGym/compare/v0.1.0...v0.1.1
107 | [0.1.0]: https://github.com/ls3-lab/QueryGym/releases/tag/v0.1.0
108 | 


--------------------------------------------------------------------------------
/examples/docker/README.md:
--------------------------------------------------------------------------------
  1 | # Docker Examples for QueryGym
  2 | 
  3 | This directory contains practical examples for using QueryGym with Docker.
  4 | 
  5 | > 📖 **Quick Reference:** See [DOCKER_SETUP.md](../../DOCKER_SETUP.md) in the root directory for a quick setup guide.
  6 | 
  7 | ## 🚀 Quick Start
  8 | 
  9 | ### Prerequisites
 10 | - Docker installed
 11 | - For GPU support: NVIDIA GPU + nvidia-container-toolkit
 12 | 
 13 | ### Setup
 14 | 
 15 | 1. **Create directories:**
 16 |    ```bash
 17 |    mkdir -p data outputs notebooks
 18 |    ```
 19 | 
 20 | 2. **Set your API key:**
 21 |    ```bash
 22 |    export OPENAI_API_KEY="sk-..."
 23 |    # Or create a .env file
 24 |    echo "OPENAI_API_KEY=sk-..." > .env
 25 |    ```
 26 | 
 27 | 3. **Run with Docker Compose:**
 28 |    ```bash
 29 |    # GPU version (interactive shell)
 30 |    docker compose run --rm querygym
 31 |    
 32 |    # CPU version (interactive shell)
 33 |    docker compose run --rm querygym-cpu
 34 |    
 35 |    # Jupyter notebook server
 36 |    docker compose up jupyter
 37 |    # Then open http://localhost:8888
 38 |    ```
 39 | 
 40 | ## 📁 What's Included
 41 | 
 42 | - **`docker-compose.yml`** - Pre-configured services using GHCR images (symlink to root)
 43 | - **`quickstart.ipynb`** - Example Jupyter notebook
 44 | - **`data/`** - Mount your datasets here
 45 | - **`outputs/`** - Reformulated queries saved here
 46 | - **`notebooks/`** - Your Jupyter notebooks
 47 | 
 48 | ## 💡 Usage Examples
 49 | 
 50 | ### Example 1: Interactive Python Session
 51 | 
 52 | ```bash
 53 | docker compose run --rm querygym
 54 | 
 55 | # Inside container:
 56 | python
 57 | >>> import querygym as qg
 58 | >>> queries = qg.load_queries("data/queries.tsv")
 59 | >>> reformulator = qg.create_reformulator("genqr")
 60 | >>> result = reformulator.reformulate(queries[0])
 61 | >>> print(result.reformulated)
 62 | ```
 63 | 
 64 | ### Example 2: Run a Script
 65 | 
 66 | ```bash
 67 | # Create your script
 68 | cat > my_experiment.py << 'EOF'
 69 | import querygym as qg
 70 | 
 71 | queries = qg.load_queries("data/queries.tsv")
 72 | reformulator = qg.create_reformulator("genqr_ensemble", model="gpt-4")
 73 | results = reformulator.reformulate_batch(queries)
 74 | 
 75 | qg.DataLoader.save_queries(
 76 |     [qg.QueryItem(r.qid, r.reformulated) for r in results],
 77 |     "outputs/reformulated.tsv"
 78 | )
 79 | print(f"Reformulated {len(results)} queries!")
 80 | EOF
 81 | 
 82 | # Run it
 83 | docker compose run --rm querygym python my_experiment.py
 84 | ```
 85 | 
 86 | ### Example 3: Jupyter Notebook
 87 | 
 88 | ```bash
 89 | # Start Jupyter server
 90 | docker compose up jupyter
 91 | 
 92 | # Open http://localhost:8888 in your browser
 93 | # Try the included quickstart.ipynb notebook
 94 | ```
 95 | 
 96 | ### Example 4: CPU-Only (No GPU)
 97 | 
 98 | ```bash
 99 | # Use the CPU version if you don't have a GPU
100 | docker compose run --rm querygym-cpu
101 | ```
102 | 
103 | ## 🔧 Configuration
104 | 
105 | ### Environment Variables
106 | 
107 | Set these in a `.env` file or export them:
108 | 
109 | ```bash
110 | OPENAI_API_KEY=sk-...              # Required for LLM-based methods
111 | OPENAI_BASE_URL=https://...        # Optional: custom API endpoint
112 | ```
113 | 
114 | ### Volume Mounts
115 | 
116 | The compose file mounts these directories:
117 | - `./data` → `/workspace/data` (your datasets)
118 | - `./outputs` → `/workspace/outputs` (results)
119 | - `./notebooks` → `/workspace/notebooks` (Jupyter notebooks)
120 | 
121 | ## 📚 More Information
122 | 
123 | - [Full Docker Guide](https://querygym.readthedocs.io/en/latest/user-guide/docker/)
124 | - [QueryGym Documentation](https://querygym.readthedocs.io/)
125 | - [GitHub Repository](https://github.com/ls3-lab/QueryGym)
126 | 
127 | ## 🐛 Troubleshooting
128 | 
129 | ### GPU Not Detected
130 | 
131 | ```bash
132 | # Verify nvidia-docker is working
133 | docker run --rm --gpus all nvidia/cuda:12.1.0-base nvidia-smi
134 | 
135 | # If not working, install nvidia-container-toolkit:
136 | # https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html
137 | ```
138 | 
139 | ### Permission Issues
140 | 
141 | ```bash
142 | # Run as your user
143 | docker compose run --rm -u $(id -u):$(id -g) querygym
144 | ```
145 | 
146 | ### Pull Latest Images
147 | 
148 | ```bash
149 | # Update to latest version
150 | docker compose pull
151 | ```
152 | 


--------------------------------------------------------------------------------
/querygym/core/searcher.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | """
  3 | Abstract Searcher Interface for querygym
  4 | 
  5 | This module defines the interface that any searcher implementation must follow
  6 | to be compatible with querygym's retrieval framework.
  7 | """
  8 | 
  9 | from abc import ABC, abstractmethod
 10 | from typing import List, Dict, Any, Optional, Union
 11 | from dataclasses import dataclass
 12 | 
 13 | 
 14 | @dataclass
 15 | class SearchHit:
 16 |     """Standardized search hit result."""
 17 |     docid: str
 18 |     score: float
 19 |     content: str
 20 |     metadata: Dict[str, Any] = None
 21 |     
 22 |     def __post_init__(self):
 23 |         if self.metadata is None:
 24 |             self.metadata = {}
 25 | 
 26 | 
 27 | class BaseSearcher(ABC):
 28 |     """
 29 |     Abstract base class for searcher implementations.
 30 |     
 31 |     Any searcher library (Pyserini, PyTerrier, etc.) can implement this interface
 32 |     to be compatible with querygym's retrieval framework.
 33 |     """
 34 |     
 35 |     @abstractmethod
 36 |     def search(self, query: str, k: int = 10, **kwargs) -> List[SearchHit]:
 37 |         """
 38 |         Search for documents using a single query.
 39 |         
 40 |         Args:
 41 |             query: The search query string
 42 |             k: Number of documents to retrieve
 43 |             **kwargs: Additional search parameters specific to the searcher
 44 |             
 45 |         Returns:
 46 |             List of SearchHit objects ordered by relevance score (highest first)
 47 |         """
 48 |         pass
 49 |     
 50 |     @abstractmethod
 51 |     def batch_search(self, queries: List[str], k: int = 10, 
 52 |                     num_threads: int = 1, **kwargs) -> List[List[SearchHit]]:
 53 |         """
 54 |         Search for documents using multiple queries in batch.
 55 |         
 56 |         Args:
 57 |             queries: List of search query strings
 58 |             k: Number of documents to retrieve per query
 59 |             num_threads: Number of threads for parallel processing
 60 |             **kwargs: Additional search parameters specific to the searcher
 61 |             
 62 |         Returns:
 63 |             List of lists of SearchHit objects, one list per query
 64 |         """
 65 |         pass
 66 |     
 67 |     @abstractmethod
 68 |     def get_searcher_info(self) -> Dict[str, Any]:
 69 |         """
 70 |         Get information about the searcher configuration.
 71 |         
 72 |         Returns:
 73 |             Dictionary containing searcher metadata (name, version, parameters, etc.)
 74 |         """
 75 |         pass
 76 |     
 77 |     def configure(self, **kwargs) -> None:
 78 |         """
 79 |         Configure searcher parameters.
 80 |         
 81 |         Args:
 82 |             **kwargs: Configuration parameters specific to the searcher
 83 |         """
 84 |         # Default implementation does nothing
 85 |         # Subclasses can override to handle specific configuration
 86 |         pass
 87 | 
 88 | 
 89 | class SearcherRegistry:
 90 |     """Registry for managing different searcher implementations."""
 91 |     
 92 |     _searchers: Dict[str, type] = {}
 93 |     
 94 |     @classmethod
 95 |     def register(cls, name: str, searcher_class: type) -> None:
 96 |         """Register a searcher implementation."""
 97 |         if not issubclass(searcher_class, BaseSearcher):
 98 |             raise ValueError(f"Searcher class must inherit from BaseSearcher")
 99 |         cls._searchers[name] = searcher_class
100 |     
101 |     @classmethod
102 |     def get_searcher(cls, name: str, **kwargs) -> BaseSearcher:
103 |         """Get a searcher instance by name."""
104 |         if name not in cls._searchers:
105 |             raise ValueError(f"Unknown searcher: {name}. Available: {list(cls._searchers.keys())}")
106 |         
107 |         searcher_class = cls._searchers[name]
108 |         return searcher_class(**kwargs)
109 |     
110 |     @classmethod
111 |     def list_searchers(cls) -> List[str]:
112 |         """List all registered searcher names."""
113 |         return list(cls._searchers.keys())
114 | 
115 | 
116 | # Convenience function for creating searchers
117 | def create_searcher(searcher_type: str, **kwargs) -> BaseSearcher:
118 |     """
119 |     Create a searcher instance.
120 |     
121 |     Args:
122 |         searcher_type: Type of searcher to create
123 |         **kwargs: Arguments to pass to the searcher constructor
124 |         
125 |     Returns:
126 |         BaseSearcher instance
127 |     """
128 |     return SearcherRegistry.get_searcher(searcher_type, **kwargs)
129 | 


--------------------------------------------------------------------------------
/docs/getting-started/quickstart.md:
--------------------------------------------------------------------------------
  1 | # Quick Start
  2 | 
  3 | This guide will help you get started with querygym in minutes.
  4 | 
  5 | ## Installation
  6 | 
  7 | ```bash
  8 | pip install querygym
  9 | ```
 10 | 
 11 | Set your OpenAI API key:
 12 | 
 13 | ```bash
 14 | export OPENAI_API_KEY="your-api-key-here"
 15 | ```
 16 | 
 17 | ## Basic Usage
 18 | 
 19 | ### 1. Load Queries
 20 | 
 21 | ```python
 22 | import querygym as qg
 23 | 
 24 | # Load from TSV file (qid<TAB>query format)
 25 | queries = qg.load_queries("queries.tsv")
 26 | 
 27 | # Or load from JSONL
 28 | queries = qg.load_queries("queries.jsonl", format="jsonl")
 29 | ```
 30 | 
 31 | ### 2. Create a Reformulator
 32 | 
 33 | ```python
 34 | # Create a reformulator with a specific method
 35 | reformulator = qg.create_reformulator(
 36 |     method_name="genqr_ensemble",
 37 |     model="gpt-4",
 38 |     params={"repeat_query_weight": 3}
 39 | )
 40 | ```
 41 | 
 42 | Available methods:
 43 | - `genqr` - Generic keyword expansion
 44 | - `genqr_ensemble` - Ensemble of keyword expansion prompts
 45 | - `query2doc` - Pseudo-document generation
 46 | - `qa_expand` - Question-answer expansion
 47 | - `mugi` - MuGI-style passage generation
 48 | - `lamer` - Context-based passage synthesis
 49 | - `query2e` - Query to entity expansion
 50 | - `csqe` - Context-based sentence extraction
 51 | 
 52 | ### 3. Reformulate Queries
 53 | 
 54 | ```python
 55 | # Reformulate a single query
 56 | query = qg.QueryItem("q1", "what causes diabetes")
 57 | result = reformulator.reformulate(query)
 58 | 
 59 | print(f"Original: {result.original}")
 60 | print(f"Reformulated: {result.reformulated}")
 61 | 
 62 | # Reformulate multiple queries
 63 | results = reformulator.reformulate_batch(queries)
 64 | ```
 65 | 
 66 | ### 4. Save Results
 67 | 
 68 | ```python
 69 | # Save reformulated queries
 70 | qg.DataLoader.save_queries(
 71 |     [qg.QueryItem(r.qid, r.reformulated) for r in results],
 72 |     "reformulated.tsv"
 73 | )
 74 | ```
 75 | 
 76 | ## Complete Example
 77 | 
 78 | ```python
 79 | import querygym as qg
 80 | 
 81 | # Load data
 82 | queries = qg.load_queries("queries.tsv")
 83 | qrels = qg.load_qrels("qrels.txt")
 84 | 
 85 | # Create reformulator
 86 | reformulator = qg.create_reformulator("genqr_ensemble", model="gpt-4")
 87 | 
 88 | # Reformulate
 89 | results = reformulator.reformulate_batch(queries)
 90 | 
 91 | # Save
 92 | qg.DataLoader.save_queries(
 93 |     [qg.QueryItem(r.qid, r.reformulated) for r in results],
 94 |     "reformulated.tsv"
 95 | )
 96 | 
 97 | print(f"Reformulated {len(results)} queries")
 98 | ```
 99 | 
100 | ## Using the CLI
101 | 
102 | querygym also provides a command-line interface:
103 | 
104 | ```bash
105 | # Run query reformulation
106 | querygym run \
107 |   --method genqr_ensemble \
108 |   --queries-tsv queries.tsv \
109 |   --output-tsv reformulated.tsv \
110 |   --cfg-path querygym/config/defaults.yaml
111 | ```
112 | 
113 | See [CLI Usage](../user-guide/cli.md) for more details.
114 | 
115 | ## Loading Datasets
116 | 
117 | ### BEIR Format
118 | 
119 | ```python
120 | import querygym as qg
121 | 
122 | # Assuming you've downloaded a BEIR dataset
123 | queries = qg.loaders.beir.load_queries("./data/nfcorpus")
124 | qrels = qg.loaders.beir.load_qrels("./data/nfcorpus", split="test")
125 | ```
126 | 
127 | ### MS MARCO Format
128 | 
129 | ```python
130 | import querygym as qg
131 | 
132 | queries = qg.loaders.msmarco.load_queries("./data/queries.tsv")
133 | qrels = qg.loaders.msmarco.load_qrels("./data/qrels.tsv")
134 | ```
135 | 
136 | See [Loading Datasets](../user-guide/datasets.md) for more details.
137 | 
138 | ## Context-Based Reformulation
139 | 
140 | Some methods (like `lamer`, `csqe`) use retrieved contexts:
141 | 
142 | ```python
143 | import querygym as qg
144 | 
145 | # Load queries and contexts
146 | queries = qg.load_queries("queries.tsv")
147 | contexts = qg.load_contexts("contexts.jsonl")
148 | 
149 | # Create context-based reformulator
150 | reformulator = qg.create_reformulator("lamer", model="gpt-4")
151 | 
152 | # Reformulate with contexts
153 | results = reformulator.reformulate_batch(queries, contexts=contexts)
154 | ```
155 | 
156 | ## Using Custom LLM Endpoints
157 | 
158 | querygym works with any OpenAI-compatible API:
159 | 
160 | ```python
161 | import os
162 | 
163 | # Set custom endpoint
164 | os.environ["OPENAI_BASE_URL"] = "https://your-endpoint.com/v1"
165 | os.environ["OPENAI_API_KEY"] = "your-key"
166 | 
167 | # Use as normal
168 | reformulator = qg.create_reformulator("genqr", model="your-model")
169 | ```
170 | 
171 | ## Next Steps
172 | 
173 | - [CLI Usage Guide](../user-guide/cli.md)
174 | - [Dataset Loading](../user-guide/datasets.md)
175 | - [Query Reformulation Methods](../user-guide/reformulation.md)
176 | - [API Reference](../api/core.md)
177 | 


--------------------------------------------------------------------------------
/examples/querygym_pyserini/reformulation_config.yaml:
--------------------------------------------------------------------------------
  1 | # QueryGym Reformulation Configuration
  2 | # 
  3 | # This file defines configurations for all reformulation methods.
  4 | # Use this file with:
  5 | #   - reformulate_queries.py: --config reformulation_config.yaml
  6 | #   - pipeline.py: --config reformulation_config.yaml
  7 | #
  8 | # Structure:
  9 | #   - Each method has its own section
 10 | #   - Global defaults can be set at the top level
 11 | #   - Method-specific configs override defaults
 12 | #   - CLI arguments override config file values
 13 | 
 14 | # Global defaults (applied to all methods unless overridden)
 15 | global:
 16 |   # LLM Configuration
 17 |   llm:
 18 |     base_url: "http://127.0.0.1:11434/v1"  # Ollama default, adjust as needed
 19 |     api_key: "ollama"  # Use "ollama" for Ollama, "EMPTY" for vLLM, or your API key
 20 |     temperature: 1.0
 21 |     max_tokens: 128
 22 |   
 23 |   # Retrieval Configuration (for context-based methods)
 24 |   retrieval:
 25 |     k: 1000  # Number of documents to retrieve per query
 26 |     threads: 16  # Number of threads for parallel retrieval
 27 |     retrieval_k: 10  # Number of documents to use as context for reformulation
 28 | 
 29 | # Method-specific configurations
 30 | methods:
 31 |   
 32 |   # GenQR: Generic Query Reformulation
 33 |   genqr:
 34 |     enabled: true
 35 |     model: "qwen2.5:7b"  # Override global model if needed
 36 |     llm:
 37 |       temperature: 0.8
 38 |       max_tokens: 256
 39 |     params:
 40 |       n_generations: 5  # Number of times to generate reformulations
 41 |   
 42 |   # GenQR Ensemble: Ensemble of multiple keyword expansion prompts
 43 |   genqr_ensemble:
 44 |     enabled: true
 45 |     model: "qwen2.5:7b"
 46 |     llm:
 47 |       temperature: 0.92  # Paper recommendation
 48 |       max_tokens: 256
 49 |     params:
 50 |       repeat_query_weight: 5  # Number of query repetitions in final output
 51 |       # variant_ids: []  # Optional: specify which variants to use (default: all 10)
 52 |       # parallel: false  # Optional: enable parallel generation (default: false)
 53 |   
 54 |   # Query2Doc: Generates pseudo-documents for the query
 55 |   query2doc:
 56 |     enabled: true
 57 |     model: "qwen2.5:7b"
 58 |     llm:
 59 |       temperature: 1.0
 60 |       max_tokens: 128  # Longer for document generation
 61 |     params: {
 62 |       "mode": "fs",
 63 |         "num_examples": 4,
 64 |         "dataset_type": "msmarco",
 65 |         "collection_path": "path/to/collection.tsv",
 66 |         "train_queries_path": "path/to/queries.train.tsv",
 67 |         "train_qrels_path": "path/to/qrels.train.tsv"
 68 |     }  # No special parameters
 69 |   
 70 |   # QA Expand: Question-answer based expansion
 71 |   qa_expand:
 72 |     enabled: true
 73 |     model: "qwen2.5:7b"
 74 |     llm:
 75 |       temperature: 1.0
 76 |       max_tokens: 256
 77 |     params: {}  # No special parameters
 78 |   
 79 |   # MuGI: Multi-granularity information expansion
 80 |   mugi:
 81 |     enabled: true
 82 |     model: "qwen2.5:7b"
 83 |     llm:
 84 |       temperature: 1.0
 85 |       max_tokens: 256
 86 |     params: {}  # No special parameters
 87 |   
 88 |   # Query2E: Query to entity expansion
 89 |   query2e:
 90 |     enabled: true
 91 |     model: "qwen2.5:7b"
 92 |     llm:
 93 |       temperature: 1.0
 94 |       max_tokens: 128
 95 |     params: {}  # No special parameters
 96 |   
 97 |   # LameR: Context-based passage synthesis (requires retrieval)
 98 |   lamer:
 99 |     enabled: true
100 |     model: "qwen2.5:7b"
101 |     llm:
102 |       temperature: 1.0
103 |       max_tokens: 128
104 |     params:
105 |       retrieval_k: 10  # Number of documents to retrieve for context
106 |       gen_passages: 5  # Number of passages to generate
107 |       # Note: Searcher is automatically configured from dataset registry
108 |   
109 |   # CSQE: Context-based sentence extraction (requires retrieval)
110 |   csqe:
111 |     enabled: true
112 |     model: "qwen2.5:7b"
113 |     llm:
114 |       temperature: 1.0
115 |       max_tokens: 1024  # Longer for sentence extraction
116 |     params:
117 |       retrieval_k: 10  # Number of documents to retrieve for context
118 |       gen_num: 2  # Number of expansions for both KEQE and CSQE (default: 2)
119 |       # Note: Searcher is automatically configured from dataset registry
120 | 
121 | # Example: Multiple configurations for the same method
122 | # You can define multiple variants with different parameters:
123 | # 
124 | # genqr_variant1:
125 | #   enabled: true
126 | #   model: "qwen2.5:7b"
127 | #   llm:
128 | #     temperature: 0.7
129 | #   params:
130 | #     n_generations: 3
131 | #
132 | # genqr_variant2:
133 | #   enabled: true
134 | #   model: "llama3.1:8b"
135 | #   llm:
136 | #     temperature: 0.9
137 | #   params:
138 | #     n_generations: 7
139 | 
140 | 


--------------------------------------------------------------------------------
/.github/CONTRIBUTING.md:
--------------------------------------------------------------------------------
  1 | # Contributing to QueryGym
  2 | 
  3 | Thank you for your interest in contributing to QueryGym! This document provides guidelines and instructions for contributing.
  4 | 
  5 | ## Table of Contents
  6 | 
  7 | - [Code of Conduct](#code-of-conduct)
  8 | - [Getting Started](#getting-started)
  9 | - [How to Contribute](#how-to-contribute)
 10 | - [Development Setup](#development-setup)
 11 | - [Pull Request Process](#pull-request-process)
 12 | - [Style Guidelines](#style-guidelines)
 13 | 
 14 | ## Code of Conduct
 15 | 
 16 | By participating in this project, you agree to maintain a respectful and inclusive environment for everyone.
 17 | 
 18 | ## Getting Started
 19 | 
 20 | 1. **Fork the repository** on GitHub
 21 | 2. **Clone your fork** locally:
 22 |    ```bash
 23 |    git clone https://github.com/YOUR_USERNAME/QueryGym.git
 24 |    cd QueryGym
 25 |    ```
 26 | 3. **Set up the development environment**:
 27 |    ```bash
 28 |    pip install -e ".[dev]"
 29 |    ```
 30 | 
 31 | ## How to Contribute
 32 | 
 33 | ### Reporting Bugs
 34 | 
 35 | - Check if the issue already exists in [GitHub Issues](https://github.com/ls3-lab/QueryGym/issues)
 36 | - If not, create a new issue using the **Bug Report** template
 37 | - Include as much detail as possible: steps to reproduce, expected vs actual behavior, environment info
 38 | 
 39 | ### Suggesting Features
 40 | 
 41 | - Open an issue using the **Feature Request** template
 42 | - Describe the feature and its use case
 43 | - Discuss the implementation approach if you have ideas
 44 | 
 45 | ### Adding a New Prompt
 46 | 
 47 | 1. Edit `querygym/prompt_bank.yaml`
 48 | 2. Add an entry with required fields:
 49 |    - `id`: Unique identifier (e.g., `method.variant.v1`)
 50 |    - `method_family`: Parent method name
 51 |    - `version`: Version string
 52 |    - `introduced_by`: Paper/source reference
 53 |    - `license`: License type
 54 |    - `authors`: List of authors
 55 |    - `tags`: Relevant tags
 56 |    - `template`: Contains `system` and `user` templates
 57 |    - `notes`: Additional notes
 58 | 
 59 | ### Adding a New Method
 60 | 
 61 | 1. Create a new file under `querygym/methods/`
 62 | 2. Subclass `BaseReformulator`
 63 | 3. Implement required methods:
 64 |    - `reformulate(self, query, contexts=None)`
 65 | 4. Register with `@register_method("method_name")`
 66 | 5. Add tests in `tests/`
 67 | 6. Update documentation
 68 | 
 69 | ## Development Setup
 70 | 
 71 | ### Prerequisites
 72 | 
 73 | - Python 3.9+
 74 | - Git
 75 | 
 76 | ### Installation
 77 | 
 78 | ```bash
 79 | # Clone your fork
 80 | git clone https://github.com/YOUR_USERNAME/QueryGym.git
 81 | cd QueryGym
 82 | 
 83 | # Create virtual environment
 84 | python -m venv venv
 85 | source venv/bin/activate  # On Windows: venv\Scripts\activate
 86 | 
 87 | # Install in development mode
 88 | pip install -e ".[dev]"
 89 | 
 90 | # Run tests
 91 | pytest
 92 | ```
 93 | 
 94 | ### Using Docker
 95 | 
 96 | ```bash
 97 | # Build development image
 98 | make build
 99 | 
100 | # Run tests in container
101 | make test
102 | ```
103 | 
104 | ## Pull Request Process
105 | 
106 | 1. **Create a branch** for your changes:
107 |    ```bash
108 |    git checkout -b feature/your-feature-name
109 |    ```
110 | 
111 | 2. **Make your changes** following the style guidelines
112 | 
113 | 3. **Add tests** for new functionality
114 | 
115 | 4. **Run tests locally**:
116 |    ```bash
117 |    pytest
118 |    ```
119 | 
120 | 5. **Commit your changes** with clear messages:
121 |    ```bash
122 |    git commit -m "Add: description of your changes"
123 |    ```
124 | 
125 | 6. **Push to your fork**:
126 |    ```bash
127 |    git push origin feature/your-feature-name
128 |    ```
129 | 
130 | 7. **Open a Pull Request** on GitHub using the PR template
131 | 
132 | 8. **Address review feedback** if requested
133 | 
134 | ### Commit Message Guidelines
135 | 
136 | - Use present tense ("Add feature" not "Added feature")
137 | - Use imperative mood ("Move cursor to..." not "Moves cursor to...")
138 | - Prefix with type: `Add:`, `Fix:`, `Update:`, `Remove:`, `Refactor:`
139 | 
140 | ## Style Guidelines
141 | 
142 | ### Python Code
143 | 
144 | - Follow [PEP 8](https://pep8.org/) style guide
145 | - Use type hints where appropriate
146 | - Write docstrings for public functions and classes
147 | - Keep functions focused and small
148 | 
149 | ### Documentation
150 | 
151 | - Update relevant documentation for any changes
152 | - Use clear, concise language
153 | - Include code examples where helpful
154 | 
155 | ### Testing
156 | 
157 | - Write tests for new features
158 | - Maintain or improve code coverage
159 | - Tests should be fast and isolated
160 | 
161 | ## Questions?
162 | 
163 | - Open a [Discussion](https://github.com/ls3-lab/QueryGym/discussions) for general questions
164 | - Check existing issues and discussions first
165 | - Join our community!
166 | 
167 | ---
168 | 
169 | Thank you for contributing to QueryGym! 🎉
170 | 


--------------------------------------------------------------------------------
/querygym/loaders/msmarco.py:
--------------------------------------------------------------------------------
  1 | """MS MARCO dataset format helpers.
  2 | 
  3 | These helpers work with MS MARCO datasets that have been downloaded locally.
  4 | They handle MS MARCO's specific TSV formats.
  5 | 
  6 | Users should download MS MARCO datasets using ir_datasets or official scripts:
  7 |     import ir_datasets
  8 |     dataset = ir_datasets.load("msmarco-passage/dev")
  9 |     # Then export to local files
 10 | 
 11 | Then use these helpers to load the data into querygym format.
 12 | """
 13 | 
 14 | from pathlib import Path
 15 | from typing import List, Dict, Union
 16 | 
 17 | from ..core.base import QueryItem
 18 | from ..data.dataloader import DataLoader
 19 | 
 20 | 
 21 | def load_queries(
 22 |     queries_tsv: Union[str, Path]
 23 | ) -> List[QueryItem]:
 24 |     """
 25 |     Load queries from MS MARCO queries TSV file.
 26 |     
 27 |     MS MARCO queries format: qid \\t query_text
 28 |     
 29 |     Args:
 30 |         queries_tsv: Path to MS MARCO queries.tsv file
 31 |         
 32 |     Returns:
 33 |         List of QueryItem objects
 34 |         
 35 |     Example:
 36 |         >>> from querygym.datasets import msmarco
 37 |         >>> queries = msmarco.load_queries("./data/msmarco_queries.tsv")
 38 |     """
 39 |     queries_tsv = Path(queries_tsv)
 40 |     
 41 |     if not queries_tsv.exists():
 42 |         raise FileNotFoundError(f"MS MARCO queries file not found: {queries_tsv}")
 43 |     
 44 |     # MS MARCO uses standard TSV: qid \t query
 45 |     return DataLoader.load_queries(
 46 |         queries_tsv,
 47 |         format="tsv",
 48 |         qid_col=0,
 49 |         query_col=1
 50 |     )
 51 | 
 52 | 
 53 | def load_qrels(
 54 |     qrels_tsv: Union[str, Path]
 55 | ) -> Dict[str, Dict[str, int]]:
 56 |     """
 57 |     Load qrels from MS MARCO qrels file.
 58 |     
 59 |     MS MARCO qrels format can be either:
 60 |     - TSV: qid \\t 0 \\t docid \\t relevance (TREC format)
 61 |     - Or: qid \\t docid \\t relevance (simplified)
 62 |     
 63 |     Args:
 64 |         qrels_tsv: Path to MS MARCO qrels file
 65 |         
 66 |     Returns:
 67 |         Dict mapping qid -> {docid -> relevance}
 68 |         
 69 |     Example:
 70 |         >>> from querygym.datasets import msmarco
 71 |         >>> qrels = msmarco.load_qrels("./data/msmarco_qrels.tsv")
 72 |     """
 73 |     qrels_tsv = Path(qrels_tsv)
 74 |     
 75 |     if not qrels_tsv.exists():
 76 |         raise FileNotFoundError(f"MS MARCO qrels file not found: {qrels_tsv}")
 77 |     
 78 |     # Try standard TREC format first
 79 |     try:
 80 |         return DataLoader.load_qrels(qrels_tsv, format="trec")
 81 |     except ValueError:
 82 |         # Fall back to simplified format (qid \t docid \t relevance)
 83 |         qrels: Dict[str, Dict[str, int]] = {}
 84 |         
 85 |         with open(qrels_tsv, "r", encoding="utf-8") as f:
 86 |             for line in f:
 87 |                 parts = line.strip().split("\t")
 88 |                 if len(parts) < 3:
 89 |                     continue
 90 |                 
 91 |                 qid = parts[0]
 92 |                 docid = parts[1]
 93 |                 try:
 94 |                     relevance = int(parts[2])
 95 |                 except ValueError:
 96 |                     continue
 97 |                 
 98 |                 if qid not in qrels:
 99 |                     qrels[qid] = {}
100 |                 qrels[qid][docid] = relevance
101 |         
102 |         if not qrels:
103 |             raise ValueError(f"No valid qrels found in {qrels_tsv}")
104 |         
105 |         return qrels
106 | 
107 | 
108 | def load_collection(
109 |     collection_tsv: Union[str, Path]
110 | ) -> Dict[str, str]:
111 |     """
112 |     Load MS MARCO passage/document collection.
113 |     
114 |     MS MARCO collection format: docid \\t text
115 |     
116 |     Args:
117 |         collection_tsv: Path to MS MARCO collection.tsv file
118 |         
119 |     Returns:
120 |         Dict mapping docid -> text
121 |         
122 |     Example:
123 |         >>> from querygym.datasets import msmarco
124 |         >>> collection = msmarco.load_collection("./data/collection.tsv")
125 |         >>> collection["doc123"]
126 |     """
127 |     collection_tsv = Path(collection_tsv)
128 |     
129 |     if not collection_tsv.exists():
130 |         raise FileNotFoundError(f"MS MARCO collection file not found: {collection_tsv}")
131 |     
132 |     collection: Dict[str, str] = {}
133 |     
134 |     with open(collection_tsv, "r", encoding="utf-8") as f:
135 |         for line in f:
136 |             parts = line.strip().split("\t")
137 |             if len(parts) < 2:
138 |                 continue
139 |             
140 |             docid = parts[0]
141 |             text = parts[1]
142 |             collection[docid] = text
143 |     
144 |     if not collection:
145 |         raise ValueError(f"No valid documents found in {collection_tsv}")
146 |     
147 |     return collection
148 | 


--------------------------------------------------------------------------------
/querygym/core/prompts.py:
--------------------------------------------------------------------------------
  1 | from __future__ import annotations
  2 | import yaml
  3 | from dataclasses import dataclass
  4 | from pathlib import Path
  5 | from typing import Dict, List, Any
  6 | 
  7 | @dataclass
  8 | class PromptSpec:
  9 |     """Specification for a single prompt template.
 10 |     
 11 |     Attributes:
 12 |         id: Unique prompt identifier
 13 |         method_family: Method family this prompt belongs to
 14 |         version: Prompt version number
 15 |         template: Template dict with 'system', 'user', 'assistant' keys
 16 |         meta: Additional metadata (authors, license, etc.)
 17 |     """
 18 |     id: str
 19 |     method_family: str
 20 |     version: int
 21 |     template: Dict[str, str]
 22 |     meta: Dict[str, Any]
 23 | 
 24 | class PromptBank:
 25 |     """Manages prompt templates from a YAML file.
 26 |     
 27 |     Loads and provides access to prompt templates with metadata.
 28 |     Supports rendering templates with variable substitution.
 29 |     
 30 |     Attributes:
 31 |         _by_id: Internal dict mapping prompt IDs to PromptSpec objects
 32 |         
 33 |     Example:
 34 |         >>> from pathlib import Path
 35 |         >>> pb = PromptBank(Path("querygym/prompt_bank.yaml"))
 36 |         >>> messages = pb.render("genqr_keywords", query="diabetes")
 37 |         >>> print(messages)
 38 |     """
 39 |     
 40 |     def __init__(self, path: str | Path):
 41 |         """Initialize the prompt bank from a YAML file.
 42 |         
 43 |         Args:
 44 |             path: Path to the prompt bank YAML file
 45 |             
 46 |         Raises:
 47 |             FileNotFoundError: If the YAML file doesn't exist
 48 |             yaml.YAMLError: If the YAML is malformed
 49 |         """
 50 |         items = yaml.safe_load(Path(path).read_text()) or []
 51 |         self._by_id: Dict[str, PromptSpec] = {}
 52 |         for x in items:
 53 |             self._by_id[x["id"]] = PromptSpec(
 54 |                 id=x["id"],
 55 |                 method_family=x.get("method_family",""),
 56 |                 version=x.get("version",1),
 57 |                 template=x["template"],
 58 |                 meta={k:v for k,v in x.items() if k not in ["id","method_family","version","template"]}
 59 |             )
 60 | 
 61 |     def render(self, prompt_id: str, **vars) -> List[Dict[str, str]]:
 62 |         """Render a prompt template with variable substitution.
 63 |         
 64 |         Args:
 65 |             prompt_id: ID of the prompt to render
 66 |             **vars: Variables to substitute in the template (e.g., query="text")
 67 |             
 68 |         Returns:
 69 |             List of message dicts with 'role' and 'content' keys,
 70 |             ready for LLM chat completion
 71 |             
 72 |         Raises:
 73 |             KeyError: If prompt_id doesn't exist
 74 |             KeyError: If required template variables are missing
 75 |             
 76 |         Example:
 77 |             >>> messages = pb.render("genqr_keywords", query="diabetes")
 78 |             >>> # Returns: [{"role": "system", "content": "..."}, ...]
 79 |         """
 80 |         spec = self._by_id[prompt_id]
 81 |         messages = []
 82 |         
 83 |         # System message (if present)
 84 |         if "system" in spec.template:
 85 |             sys = spec.template.get("system","").format(**vars)
 86 |             messages.append({"role":"system","content":sys})
 87 |         
 88 |         # User message (if present)
 89 |         if "user" in spec.template:
 90 |             usr = spec.template.get("user","").format(**vars)
 91 |             messages.append({"role":"user","content":usr})
 92 |         
 93 |         # Assistant message (if present) - for priming responses
 94 |         if "assistant" in spec.template:
 95 |             asst = spec.template.get("assistant","").format(**vars)
 96 |             messages.append({"role":"assistant","content":asst})
 97 |         
 98 |         return messages
 99 | 
100 |     def list(self) -> List[str]:
101 |         """List all available prompt IDs.
102 |         
103 |         Returns:
104 |             List of prompt IDs
105 |         """
106 |         return list(self._by_id.keys())
107 | 
108 |     def get(self, prompt_id: str) -> PromptSpec:
109 |         """Get a prompt specification by ID.
110 |         
111 |         Args:
112 |             prompt_id: ID of the prompt
113 |             
114 |         Returns:
115 |             PromptSpec object
116 |             
117 |         Raises:
118 |             KeyError: If prompt_id doesn't exist
119 |         """
120 |         return self._by_id[prompt_id]
121 |     
122 |     def get_meta(self, prompt_id: str) -> Dict[str, Any]:
123 |         """Get metadata for a prompt.
124 |         
125 |         Args:
126 |             prompt_id: ID of the prompt
127 |             
128 |         Returns:
129 |             Metadata dict (authors, license, etc.)
130 |             
131 |         Raises:
132 |             KeyError: If prompt_id doesn't exist
133 |         """
134 |         return self._by_id[prompt_id].meta
135 | 


--------------------------------------------------------------------------------
/querygym/methods/lamer.py:
--------------------------------------------------------------------------------
  1 | from __future__ import annotations
  2 | from ..core.base import BaseReformulator, QueryItem, ReformulationResult
  3 | from ..core.registry import register_method
  4 | from typing import List, Optional, Dict, Any
  5 | 
  6 | @register_method("lamer")
  7 | class LameR(BaseReformulator):
  8 |     VERSION = "0.1"
  9 |     REQUIRES_CONTEXT = True
 10 |     CONCATENATION_STRATEGY = "interleaved_query_content"  # q + a1 + q + a2 + q + a3 + q + a4 + q + a5
 11 | 
 12 |     def reformulate(self, q: QueryItem, contexts=None) -> ReformulationResult:
 13 |         # Use provided contexts (from batch retrieval)
 14 |         ctxs = contexts or []
 15 | 
 16 |         # Prepare contexts for prompting (use top-10 by default)
 17 |         retrieval_k = int(self.cfg.params.get("retrieval_k", 10))
 18 |         prompt_ctxs = ctxs[:retrieval_k]
 19 |         contexts_blob = "\n".join([f"{i+1}. {psg}" for i, psg in enumerate(prompt_ctxs)])
 20 | 
 21 |         # Prompt LLM to generate multiple passages given contexts
 22 |         gen_passages: List[str] = []
 23 |         num_generations = int(self.cfg.params.get("gen_passages", 5))
 24 |         max_tokens = int(self.cfg.llm.get("max_tokens", 128))
 25 |         temperature = float(self.cfg.llm.get("temperature", 1.0))
 26 | 
 27 |         # Use prompt template defined in prompt_bank.yaml (lamer.msmarco.v1)
 28 |         msgs = self.prompts.render("lamer.msmarco.v1", query=q.text, contexts=contexts_blob)
 29 |         for _ in range(num_generations):
 30 |             passage = self.llm.chat(msgs, temperature=temperature, max_tokens=max_tokens)
 31 |             # Clean up quotes from the generated passage
 32 |             passage = passage.strip().strip('"').strip("'")
 33 |             gen_passages.append(passage)
 34 | 
 35 |         # Construct final expanded query using interleaved_query_content
 36 |         reformulated = self.concatenate_result(q.text, gen_passages)
 37 | 
 38 |         return ReformulationResult(
 39 |             q.qid,
 40 |             q.text,
 41 |             reformulated,
 42 |             metadata={
 43 |                 "generated_passages": gen_passages,  # Store the actual passages
 44 |                 "generated_passages_count": len(gen_passages),
 45 |                 "used_ctx": len(prompt_ctxs),
 46 |             },
 47 |         )
 48 | 
 49 |     def _get_retrieval_params(self) -> Optional[Dict[str, Any]]:
 50 |         """Get LameR-specific retrieval parameters for batch retrieval.
 51 |         
 52 |         Returns retrieval parameters compatible with BaseSearcher interface.
 53 |         Supports both searcher_type/searcher_kwargs format and legacy format.
 54 |         """
 55 |         # Check if searcher is provided directly (wrappers, adapter instances, etc.)
 56 |         if "searcher" in self.cfg.params:
 57 |             # Direct searcher instance provided
 58 |             return {
 59 |                 "searcher": self.cfg.params["searcher"],
 60 |                 "k": int(self.cfg.params.get("retrieval_k", 10)),
 61 |                 "threads": int(self.cfg.params.get("threads", 16)),
 62 |             }
 63 |         
 64 |         # Check if using new searcher interface format with searcher_type
 65 |         if "searcher_type" in self.cfg.params:
 66 |             # New format: use searcher_type and searcher_kwargs
 67 |             searcher_type = self.cfg.params.get("searcher_type", "pyserini")
 68 |             searcher_kwargs = self.cfg.params.get("searcher_kwargs", {})
 69 |             
 70 |             # If index is provided in params but not in searcher_kwargs, add it
 71 |             if "index" in self.cfg.params and "index" not in searcher_kwargs:
 72 |                 searcher_kwargs["index"] = self.cfg.params["index"]
 73 |             
 74 |             return {
 75 |                 "searcher_type": searcher_type,
 76 |                 "searcher_kwargs": searcher_kwargs,
 77 |                 "k": int(self.cfg.params.get("retrieval_k", 10)),
 78 |                 "threads": int(self.cfg.params.get("threads", 16)),
 79 |             }
 80 |         
 81 |         # Legacy format: convert old-style params to new format
 82 |         index = self.cfg.params.get("index")
 83 |         if not index:
 84 |             return None
 85 |         
 86 |         # Build searcher_kwargs from legacy params
 87 |         searcher_kwargs = {
 88 |             "index": index,
 89 |             "searcher_type": "impact" if self.cfg.params.get("impact", False) else "bm25",
 90 |             "answer_key": self.cfg.params.get("answer_key", "contents"),
 91 |         }
 92 |         
 93 |         # Add BM25 parameters if specified
 94 |         k1 = self.cfg.params.get("k1")
 95 |         b = self.cfg.params.get("b")
 96 |         if k1 is not None and b is not None:
 97 |             searcher_kwargs["k1"] = k1
 98 |             searcher_kwargs["b"] = b
 99 |         
100 |         # Add RM3 if requested
101 |         if self.cfg.params.get("rm3", False):
102 |             searcher_kwargs["rm3"] = True
103 |         
104 |         # Add Rocchio if requested
105 |         if self.cfg.params.get("rocchio", False):
106 |             searcher_kwargs["rocchio"] = True
107 |             if self.cfg.params.get("rocchio_use_negative", False):
108 |                 searcher_kwargs["rocchio_use_negative"] = True
109 |         
110 |         return {
111 |             "searcher_type": "pyserini",  # Default to pyserini for legacy format
112 |             "searcher_kwargs": searcher_kwargs,
113 |             "k": int(self.cfg.params.get("retrieval_k", 10)),
114 |             "threads": int(self.cfg.params.get("threads", 16)),
115 |         }
116 | 


--------------------------------------------------------------------------------
/querygym/loaders/beir.py:
--------------------------------------------------------------------------------
  1 | """BEIR dataset format helpers.
  2 | 
  3 | These helpers work with BEIR datasets that have been downloaded locally.
  4 | They handle BEIR's specific file formats and directory structures.
  5 | 
  6 | Users should download BEIR datasets using the official BEIR library:
  7 |     from beir.datasets.data_loader import GenericDataLoader
  8 |     data_path = GenericDataLoader("nfcorpus").download_and_unzip()
  9 | 
 10 | Then use these helpers to load the data into querygym format.
 11 | """
 12 | 
 13 | from pathlib import Path
 14 | from typing import List, Dict, Union
 15 | import json
 16 | 
 17 | from ..core.base import QueryItem
 18 | from ..data.dataloader import DataLoader
 19 | 
 20 | 
 21 | def load_queries(
 22 |     beir_data_dir: Union[str, Path],
 23 |     split: str = "test"
 24 | ) -> List[QueryItem]:
 25 |     """
 26 |     Load queries from a BEIR dataset directory.
 27 |     
 28 |     BEIR datasets have a queries.jsonl file with format:
 29 |         {"_id": "query_id", "text": "query text", ...}
 30 |     
 31 |     Args:
 32 |         beir_data_dir: Path to BEIR dataset directory (contains queries.jsonl)
 33 |         split: Dataset split (not used for queries, kept for API consistency)
 34 |         
 35 |     Returns:
 36 |         List of QueryItem objects
 37 |         
 38 |     Example:
 39 |         >>> from querygym.datasets import beir
 40 |         >>> queries = beir.load_queries("./data/nfcorpus")
 41 |     """
 42 |     beir_data_dir = Path(beir_data_dir)
 43 |     queries_file = beir_data_dir / "queries.jsonl"
 44 |     
 45 |     if not queries_file.exists():
 46 |         raise FileNotFoundError(
 47 |             f"BEIR queries file not found: {queries_file}\n"
 48 |             f"Expected BEIR directory structure with queries.jsonl"
 49 |         )
 50 |     
 51 |     # BEIR uses "_id" and "text" as keys
 52 |     return DataLoader.load_queries(
 53 |         queries_file,
 54 |         format="jsonl",
 55 |         qid_key="_id",
 56 |         query_key="text"
 57 |     )
 58 | 
 59 | 
 60 | def load_qrels(
 61 |     beir_data_dir: Union[str, Path],
 62 |     split: str = "test"
 63 | ) -> Dict[str, Dict[str, int]]:
 64 |     """
 65 |     Load qrels from a BEIR dataset directory.
 66 |     
 67 |     BEIR qrels are in TSV format: query-id corpus-id score
 68 |     Located at: qrels/{split}.tsv
 69 |     
 70 |     Args:
 71 |         beir_data_dir: Path to BEIR dataset directory
 72 |         split: Dataset split ("train", "dev", "test")
 73 |         
 74 |     Returns:
 75 |         Dict mapping qid -> {docid -> relevance}
 76 |         
 77 |     Example:
 78 |         >>> from querygym.datasets import beir
 79 |         >>> qrels = beir.load_qrels("./data/nfcorpus", split="test")
 80 |     """
 81 |     beir_data_dir = Path(beir_data_dir)
 82 |     qrels_file = beir_data_dir / "qrels" / f"{split}.tsv"
 83 |     
 84 |     if not qrels_file.exists():
 85 |         raise FileNotFoundError(
 86 |             f"BEIR qrels file not found: {qrels_file}\n"
 87 |             f"Expected: {beir_data_dir}/qrels/{split}.tsv"
 88 |         )
 89 |     
 90 |     # BEIR qrels format: query-id \t corpus-id \t score
 91 |     # We need to convert to standard TREC format (add iteration column)
 92 |     qrels: Dict[str, Dict[str, int]] = {}
 93 |     
 94 |     with open(qrels_file, "r", encoding="utf-8") as f:
 95 |         for line in f:
 96 |             parts = line.strip().split("\t")
 97 |             if len(parts) < 3:
 98 |                 continue
 99 |             
100 |             qid = parts[0]
101 |             docid = parts[1]
102 |             try:
103 |                 relevance = int(parts[2])
104 |             except ValueError:
105 |                 continue
106 |             
107 |             if qid not in qrels:
108 |                 qrels[qid] = {}
109 |             qrels[qid][docid] = relevance
110 |     
111 |     if not qrels:
112 |         raise ValueError(f"No valid qrels found in {qrels_file}")
113 |     
114 |     return qrels
115 | 
116 | 
117 | def load_corpus(
118 |     beir_data_dir: Union[str, Path]
119 | ) -> Dict[str, Dict[str, str]]:
120 |     """
121 |     Load corpus from a BEIR dataset directory.
122 |     
123 |     BEIR corpus is in JSONL format with:
124 |         {"_id": "doc_id", "title": "...", "text": "...", ...}
125 |     
126 |     Args:
127 |         beir_data_dir: Path to BEIR dataset directory
128 |         
129 |     Returns:
130 |         Dict mapping doc_id -> {"title": ..., "text": ...}
131 |         
132 |     Example:
133 |         >>> from querygym.datasets import beir
134 |         >>> corpus = beir.load_corpus("./data/nfcorpus")
135 |         >>> corpus["doc123"]["title"]
136 |     """
137 |     beir_data_dir = Path(beir_data_dir)
138 |     corpus_file = beir_data_dir / "corpus.jsonl"
139 |     
140 |     if not corpus_file.exists():
141 |         raise FileNotFoundError(
142 |             f"BEIR corpus file not found: {corpus_file}\n"
143 |             f"Expected BEIR directory structure with corpus.jsonl"
144 |         )
145 |     
146 |     corpus: Dict[str, Dict[str, str]] = {}
147 |     
148 |     with open(corpus_file, "r", encoding="utf-8") as f:
149 |         for line in f:
150 |             if not line.strip():
151 |                 continue
152 |             
153 |             try:
154 |                 doc = json.loads(line)
155 |             except json.JSONDecodeError:
156 |                 continue
157 |             
158 |             doc_id = doc.get("_id")
159 |             if not doc_id:
160 |                 continue
161 |             
162 |             corpus[str(doc_id)] = {
163 |                 "title": doc.get("title", ""),
164 |                 "text": doc.get("text", "")
165 |             }
166 |     
167 |     if not corpus:
168 |         raise ValueError(f"No valid documents found in {corpus_file}")
169 |     
170 |     return corpus
171 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | [![Publish to PyPI](https://github.com/ls3-lab/QueryGym/actions/workflows/publish.yml/badge.svg)](https://github.com/ls3-lab/QueryGym/actions/workflows/publish.yml)
  2 | [![Build and Push Docker Images](https://github.com/ls3-lab/QueryGym/actions/workflows/docker.yml/badge.svg)](https://github.com/ls3-lab/QueryGym/actions/workflows/docker.yml)
  3 | [![PyPI version](https://badge.fury.io/py/querygym.svg)](https://pypi.org/project/querygym/)
  4 | ![PyPI - Downloads](https://img.shields.io/pypi/dw/querygym?color=blueviolet&label=downloads)
  5 | [![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
  6 | [![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
  7 | 
  8 | 
  9 | <div align="center">
 10 |   <img src="https://raw.githubusercontent.com/ls3-lab/QueryGym/main/docs/querygym-logo.png" alt="QueryGym Logo" width="600">
 11 | </div>
 12 | 
 13 | <p align="center">
 14 |   <strong>A lightweight, reproducible toolkit for LLM-based query reformulation</strong>
 15 | </p>
 16 | 
 17 | <p align="center">
 18 |   <a href="https://querygym.readthedocs.io/">📚 Documentation</a> •
 19 |   <a href="https://ls3-lab.github.io/QueryGym/leaderboard.html">📊 Leaderboard</a> •
 20 |   <a href="https://pypi.org/project/querygym/">📦 PyPI</a> •
 21 |   <a href="https://arxiv.org/abs/2511.15996">📄 Paper</a>
 22 | </p>
 23 | 
 24 | ---
 25 | 
 26 | ## Features
 27 | 
 28 | - Single **Prompt Bank** (YAML) with metadata
 29 | - **Simple DataLoader**: Dependency-free file loading for queries, qrels, and contexts
 30 | - **Format Loaders**: Optional BEIR and MS MARCO format loaders in `querygym.loaders`
 31 | - **OpenAI-compatible** LLM client (works with any OpenAI API–compatible endpoint)
 32 | - **Pyserini** optional: either pass contexts (JSONL) or pass a retriever instance to build contexts
 33 | - Export-only: emits reformulated queries; optionally generates a **bash** script for Pyserini + `trec_eval`
 34 | 
 35 | ## Installation
 36 | 
 37 | ### Option 1: Install from PyPI
 38 | ```bash
 39 | pip install querygym
 40 | ```
 41 | 
 42 | ### Option 2: Use Docker (Recommended for Quick Start)
 43 | ```bash
 44 | # GPU version (default)
 45 | docker pull ghcr.io/ls3-lab/querygym:latest
 46 | docker run -it --gpus all ghcr.io/ls3-lab/querygym:latest
 47 | 
 48 | # CPU version (lightweight)
 49 | docker pull ghcr.io/ls3-lab/querygym:cpu
 50 | docker run -it ghcr.io/ls3-lab/querygym:cpu
 51 | 
 52 | # Or use Docker Compose
 53 | docker compose run --rm querygym
 54 | ```
 55 | 
 56 | 📖 **Docker Setup:** See [DOCKER_SETUP.md](DOCKER_SETUP.md) for quick start or the [full Docker guide](https://querygym.readthedocs.io/en/latest/user-guide/docker/) for detailed usage.
 57 | 
 58 | ## Quickstart
 59 | 
 60 | ### Python API (Recommended)
 61 | ```python
 62 | import querygym as qg
 63 | 
 64 | # Load data
 65 | queries = qg.load_queries("queries.tsv")
 66 | qrels = qg.load_qrels("qrels.txt")
 67 | contexts = qg.load_contexts("contexts.jsonl")
 68 | 
 69 | # Create reformulator
 70 | reformulator = qg.create_reformulator("genqr_ensemble", model="gpt-4")
 71 | 
 72 | # Reformulate
 73 | results = reformulator.reformulate_batch(queries)
 74 | 
 75 | # Save
 76 | qg.DataLoader.save_queries(
 77 |     [qg.QueryItem(r.qid, r.reformulated) for r in results],
 78 |     "reformulated.tsv"
 79 | )
 80 | ```
 81 | 
 82 | ### CLI
 83 | ```bash
 84 | pip install -e .[hf,beir,dev]
 85 | export OPENAI_API_KEY=sk-...
 86 | 
 87 | # Run a method (e.g., genqr_ensemble)
 88 | querygym run --method genqr_ensemble \
 89 |   --queries-tsv queries.tsv \
 90 |   --output-tsv reformulated.tsv \
 91 |   --cfg-path querygym/config/defaults.yaml
 92 | ```
 93 | 
 94 | ### Loading Datasets
 95 | 
 96 | **BEIR:**
 97 | ```python
 98 | import querygym as qg
 99 | 
100 | # Download with BEIR library
101 | from beir.datasets.data_loader import GenericDataLoader
102 | data_path = GenericDataLoader("nfcorpus").download_and_unzip()
103 | 
104 | # Load with querygym
105 | queries = qg.loaders.beir.load_queries(data_path)
106 | qrels = qg.loaders.beir.load_qrels(data_path)
107 | ```
108 | 
109 | **MS MARCO:**
110 | ```python
111 | import querygym as qg
112 | 
113 | # Load from local files (download with ir_datasets)
114 | queries = qg.loaders.msmarco.load_queries("queries.tsv")
115 | qrels = qg.loaders.msmarco.load_qrels("qrels.tsv")
116 | ```
117 | 
118 | ## Examples
119 | 
120 | See the [examples](examples/) directory for:
121 | - **[Code snippets](examples/snippets/)** - Quick reference examples
122 | - **[Docker examples](examples/docker/)** - Containerized workflows with Jupyter notebooks
123 | - **[QueryGym + Pyserini](examples/querygym_pyserini/)** - Complete retrieval pipelines
124 | - **[Methods Reference](docs/user-guide/methods-reference.md)** - Complete guide to all query reformulation methods
125 | 
126 | Check [examples/README.md](examples/README.md) for the full guide.
127 | 
128 | ## Contributing
129 | 
130 | We welcome contributions! Here's how you can help:
131 | 
132 | ### Adding a New Prompt
133 | 1. Edit `querygym/prompt_bank.yaml`
134 | 2. Add an entry with fields: `id`, `method_family`, `version`, `introduced_by`, `license`, `authors`, `tags`, `template:{system,user}`, `notes`
135 | 
136 | ### Adding a New Method
137 | 1. Create a class under `querygym/methods/*.py`
138 | 2. Subclass `BaseReformulator`, annotate `VERSION`, and register with `@register_method("name")`
139 | 3. Pull templates via `PromptBank.render(prompt_id, query=...)`
140 | 
141 | ### Reporting Issues
142 | - Found a bug? [Open an issue](https://github.com/ls3-lab/QueryGym/issues)
143 | - Have a feature request? We'd love to hear it!
144 | 
145 | For detailed development guidelines, see the [Contributing Guide](https://querygym.readthedocs.io/en/latest/development/contributing/) in our documentation.
146 | 
147 | ## License
148 | 
149 | This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENSE) file for details.
150 | 
151 | ### Citation
152 | 
153 | If you use QueryGym in your research, please cite:
154 | 
155 | ```bibtex
156 | @misc{bigdeli2025querygymtoolkitreproduciblellmbased,
157 |       title={QueryGym: A Toolkit for Reproducible LLM-Based Query Reformulation}, 
158 |       author={Amin Bigdeli and Radin Hamidi Rad and Mert Incesu and Negar Arabzadeh and Charles L. A. Clarke and Ebrahim Bagheri},
159 |       year={2025},
160 |       eprint={2511.15996},
161 |       archivePrefix={arXiv},
162 |       primaryClass={cs.IR},
163 |       url={https://arxiv.org/abs/2511.15996}, 
164 | }
165 | ```
166 | 


--------------------------------------------------------------------------------
/querygym/methods/mugi.py:
--------------------------------------------------------------------------------
  1 | from __future__ import annotations
  2 | from typing import List
  3 | from ..core.base import BaseReformulator, QueryItem, ReformulationResult
  4 | from ..core.registry import register_method
  5 | 
  6 | @register_method("mugi")
  7 | class MuGI(BaseReformulator):
  8 |     """
  9 |     MuGI (Multi-Text Generation Integration) Method
 10 |     
 11 |     Generates multiple diverse pseudo-documents per query and uses adaptive
 12 |     concatenation to balance term frequencies between short queries and long pseudo-docs.
 13 |     
 14 |     Key Parameters:
 15 |         - num_docs: Number of pseudo-documents to generate per query (default: 5)
 16 |         - adaptive_times: Divisor for adaptive repetition ratio (default: 5)
 17 |         - max_tokens: Maximum tokens per pseudo-document (default: 1024)
 18 |         - temperature: Sampling temperature for diversity (default: 1.0)
 19 |         - mode: "zs" for zero-shot or "fs" for few-shot (default: "zs")
 20 |         - prompt_id: (Advanced) Direct prompt ID override ("mugi.zeroshot.v1" or "mugi.fewshot.v1")
 21 |         - parallel: If True, generate all pseudo-docs in parallel; if False, sequential (default: False)
 22 |     
 23 |     Formula:
 24 |         repetition_times = (len(all_pseudo_docs) // len(query)) // adaptive_times
 25 |         enhanced_query = (query + ' ') * repetition_times + all_pseudo_docs
 26 |     """
 27 |     VERSION = "1.0"
 28 |     REQUIRES_CONTEXT = False
 29 |     CONCATENATION_STRATEGY = "adaptive_query_repeat_plus_generated"
 30 | 
 31 |     def reformulate(self, q: QueryItem, contexts=None) -> ReformulationResult:
 32 |         # Get configurable parameters from config
 33 |         num_docs = int(self.cfg.params.get("num_docs", 5))
 34 |         adaptive_times = int(self.cfg.params.get("adaptive_times", 6))
 35 |         max_tokens = int(self.cfg.params.get("max_tokens", self.cfg.llm.get("max_tokens", 1024)))
 36 |         temperature = float(self.cfg.params.get("temperature", self.cfg.llm.get("temperature", 1.0)))
 37 |         parallel = bool(self.cfg.params.get("parallel", False))
 38 |         
 39 |         # Mode selection: "zs" (zero-shot) or "fs" (few-shot)
 40 |         mode = str(self.cfg.params.get("mode", "zs"))
 41 |         
 42 |         # Map mode to prompt_id (or allow direct prompt_id override for advanced users)
 43 |         if "prompt_id" in self.cfg.params:
 44 |             # Advanced: allow direct prompt ID specification
 45 |             prompt_id = str(self.cfg.params.get("prompt_id"))
 46 |         else:
 47 |             # Simple: use mode to select prompt
 48 |             if mode in ["fs", "fewshot"]:
 49 |                 prompt_id = "mugi.fewshot.v1"
 50 |             elif mode in ["zs", "zeroshot"]:
 51 |                 prompt_id = "mugi.zeroshot.v1"
 52 |             else:
 53 |                 raise ValueError(f"Invalid mode '{mode}' for MuGI. Must be 'zs' (zero-shot) or 'fs' (few-shot).")
 54 |         
 55 |         # Generate multiple pseudo-documents for diversity and coverage
 56 |         pseudo_docs: List[str] = []
 57 |         
 58 |         if parallel:
 59 |             # Parallel generation using ThreadPoolExecutor
 60 |             from concurrent.futures import ThreadPoolExecutor, as_completed
 61 |             
 62 |             def generate_single_doc(doc_idx: int) -> str:
 63 |                 msgs = self.prompts.render(prompt_id, query=q.text)
 64 |                 pseudo_doc = self.llm.chat(msgs, temperature=temperature, max_tokens=max_tokens)
 65 |                 return pseudo_doc.strip().strip('"').strip("'")
 66 |             
 67 |             # Generate all pseudo-docs in parallel
 68 |             with ThreadPoolExecutor(max_workers=num_docs) as executor:
 69 |                 futures = [executor.submit(generate_single_doc, i) for i in range(num_docs)]
 70 |                 for future in as_completed(futures):
 71 |                     pseudo_docs.append(future.result())
 72 |         else:
 73 |             # Sequential generation (original behavior)
 74 |             for i in range(num_docs):
 75 |                 # Render prompt (mugi.zeroshot.v1 or mugi.fewshot.v1)
 76 |                 msgs = self.prompts.render(prompt_id, query=q.text)
 77 |                 
 78 |                 # Generate pseudo-document with high temperature for diversity
 79 |                 pseudo_doc = self.llm.chat(
 80 |                     msgs, 
 81 |                     temperature=temperature,
 82 |                     max_tokens=max_tokens
 83 |                 )
 84 |                 
 85 |                 # Clean up the generated pseudo-document
 86 |                 pseudo_doc = pseudo_doc.strip().strip('"').strip("'")
 87 |                 pseudo_docs.append(pseudo_doc)
 88 |         
 89 |         # Join all pseudo-documents into single text
 90 |         all_pseudo_docs = ' '.join(pseudo_docs)
 91 |         
 92 |         # Use base class concatenation with adaptive_query_repeat_plus_generated strategy
 93 |         # This automatically handles the MuGI formula and cleaning
 94 |         reformulated = self.concatenate_result(q.text, all_pseudo_docs)
 95 |         
 96 |         # Calculate metrics for metadata (after concatenation)
 97 |         query_len = len(q.text)
 98 |         docs_len = len(all_pseudo_docs)
 99 |         if query_len > 0:
100 |             repetition_times = max(1, (docs_len // query_len) // adaptive_times)
101 |         else:
102 |             repetition_times = 1
103 |         
104 |         # Return result with metadata
105 |         return ReformulationResult(
106 |             q.qid,
107 |             q.text,
108 |             reformulated,
109 |             metadata={
110 |                 "pseudo_docs": pseudo_docs,
111 |                 "num_docs": num_docs,
112 |                 "adaptive_times": adaptive_times,
113 |                 "repetition_times": repetition_times,
114 |                 "query_len": query_len,
115 |                 "docs_len": docs_len,
116 |                 "mode": mode,
117 |                 "prompt_id": prompt_id,
118 |                 "parallel": parallel,
119 |                 # Store individual pseudo-docs for analysis
120 |                 "pseudo_doc_1": pseudo_docs[0] if len(pseudo_docs) > 0 else "",
121 |                 "pseudo_doc_2": pseudo_docs[1] if len(pseudo_docs) > 1 else "",
122 |                 "pseudo_doc_3": pseudo_docs[2] if len(pseudo_docs) > 2 else "",
123 |                 "pseudo_doc_4": pseudo_docs[3] if len(pseudo_docs) > 3 else "",
124 |                 "pseudo_doc_5": pseudo_docs[4] if len(pseudo_docs) > 4 else "",
125 |             }
126 |         )
127 | 


--------------------------------------------------------------------------------
/docs/development/contributing.md:
--------------------------------------------------------------------------------
  1 | # Contributing to querygym
  2 | 
  3 | Thank you for your interest in contributing to querygym! This guide will help you get started.
  4 | 
  5 | ## Getting Started
  6 | 
  7 | ### 1. Fork and Clone
  8 | 
  9 | ```bash
 10 | # Fork the repository on GitHub
 11 | # Then clone your fork
 12 | git clone https://github.com/YOUR_USERNAME/queryGym.git
 13 | cd queryGym
 14 | ```
 15 | 
 16 | ### 2. Set Up Development Environment
 17 | 
 18 | ```bash
 19 | # Create a virtual environment
 20 | python -m venv venv
 21 | source venv/bin/activate  # On Windows: venv\Scripts\activate
 22 | 
 23 | # Install in editable mode with dev dependencies
 24 | pip install -e .[dev]
 25 | ```
 26 | 
 27 | ### 3. Create a Branch
 28 | 
 29 | ```bash
 30 | git checkout -b feature/your-feature-name
 31 | ```
 32 | 
 33 | ## Development Workflow
 34 | 
 35 | ### Running Tests
 36 | 
 37 | ```bash
 38 | # Run all tests
 39 | pytest
 40 | 
 41 | # Run with coverage
 42 | pytest --cov=querygym
 43 | 
 44 | # Run specific test file
 45 | pytest tests/test_prompts.py
 46 | 
 47 | # Run with verbose output
 48 | pytest -v
 49 | ```
 50 | 
 51 | ### Code Quality
 52 | 
 53 | We use several tools to maintain code quality:
 54 | 
 55 | ```bash
 56 | # Format code with black
 57 | black querygym/ tests/
 58 | 
 59 | # Lint with ruff
 60 | ruff check querygym/ tests/
 61 | 
 62 | # Type checking with mypy
 63 | mypy querygym/
 64 | ```
 65 | 
 66 | ### Pre-commit Checks
 67 | 
 68 | Before committing, ensure:
 69 | 
 70 | 1. All tests pass: `pytest`
 71 | 2. Code is formatted: `black querygym/ tests/`
 72 | 3. No linting errors: `ruff check querygym/ tests/`
 73 | 4. Type hints are correct: `mypy querygym/`
 74 | 
 75 | ## Adding a New Reformulation Method
 76 | 
 77 | ### 1. Create Method File
 78 | 
 79 | Create a new file in `querygym/methods/`:
 80 | 
 81 | ```python
 82 | # querygym/methods/my_method.py
 83 | from __future__ import annotations
 84 | from typing import Optional, List
 85 | from ..core.base import BaseReformulator, QueryItem, ReformulationResult, MethodConfig
 86 | 
 87 | class MyMethod(BaseReformulator):
 88 |     VERSION = "1.0"
 89 |     REQUIRES_CONTEXT = False  # Set to True if method needs contexts
 90 |     
 91 |     def reformulate(self, q: QueryItem, contexts: Optional[List[str]] = None) -> ReformulationResult:
 92 |         # Get prompt
 93 |         prompt = self.prompts.get("my_method_prompt_id")
 94 |         
 95 |         # Build messages
 96 |         messages = [
 97 |             {"role": "system", "content": prompt.template["system"]},
 98 |             {"role": "user", "content": prompt.template["user"].format(query=q.text)}
 99 |         ]
100 |         
101 |         # Call LLM
102 |         response = self.llm.chat(messages)
103 |         
104 |         # Concatenate with original query
105 |         reformulated = self.concatenate_result(q.text, response)
106 |         
107 |         return ReformulationResult(
108 |             qid=q.qid,
109 |             original=q.text,
110 |             reformulated=reformulated
111 |         )
112 | ```
113 | 
114 | ### 2. Add Prompt to Prompt Bank
115 | 
116 | Add your prompt to `querygym/prompt_bank.yaml`:
117 | 
118 | ```yaml
119 | my_method_prompt_id:
120 |   meta:
121 |     version: 1
122 |     introduced_by: "Your Name"
123 |     license: "CC-BY-4.0"
124 |     authors: ["Your Name"]
125 |   method_family: "my_method"
126 |   template:
127 |     system: "You are a helpful assistant for query reformulation."
128 |     user: "Reformulate this query: {query}"
129 | ```
130 | 
131 | ### 3. Register the Method
132 | 
133 | Add to `querygym/methods/__init__.py`:
134 | 
135 | ```python
136 | from .my_method import MyMethod
137 | 
138 | __all__ = [
139 |     # ... existing methods
140 |     "MyMethod",
141 | ]
142 | ```
143 | 
144 | Register in `querygym/core/registry.py`:
145 | 
146 | ```python
147 | from ..methods.my_method import MyMethod
148 | 
149 | register_method("my_method", MyMethod)
150 | ```
151 | 
152 | ### 4. Add Tests
153 | 
154 | Create `tests/test_my_method.py`:
155 | 
156 | ```python
157 | from querygym.methods.my_method import MyMethod
158 | from querygym.core.base import QueryItem, MethodConfig
159 | from pathlib import Path
160 | 
161 | def test_my_method():
162 |     # Setup
163 |     cfg = MethodConfig(
164 |         name="my_method",
165 |         params={},
166 |         llm={"model": "gpt-4"}
167 |     )
168 |     
169 |     # Mock LLM client
170 |     class MockLLM:
171 |         def chat(self, messages):
172 |             return "expanded query"
173 |     
174 |     # Mock prompt resolver
175 |     class MockPrompts:
176 |         def get(self, prompt_id):
177 |             class Prompt:
178 |                 template = {
179 |                     "system": "System prompt",
180 |                     "user": "User: {query}"
181 |                 }
182 |             return Prompt()
183 |     
184 |     method = MyMethod(cfg, MockLLM(), MockPrompts())
185 |     
186 |     # Test
187 |     query = QueryItem("q1", "test query")
188 |     result = method.reformulate(query)
189 |     
190 |     assert result.qid == "q1"
191 |     assert result.original == "test query"
192 |     assert "expanded query" in result.reformulated
193 | ```
194 | 
195 | ### 5. Update Documentation
196 | 
197 | Add documentation in `docs/user-guide/reformulation.md`.
198 | 
199 | ## Pull Request Process
200 | 
201 | 1. **Update your branch**:
202 |    ```bash
203 |    git fetch upstream
204 |    git rebase upstream/main
205 |    ```
206 | 
207 | 2. **Run all checks**:
208 |    ```bash
209 |    pytest
210 |    black querygym/ tests/
211 |    ruff check querygym/ tests/
212 |    mypy querygym/
213 |    ```
214 | 
215 | 3. **Commit your changes**:
216 |    ```bash
217 |    git add .
218 |    git commit -m "Add: Brief description of changes"
219 |    ```
220 | 
221 | 4. **Push to your fork**:
222 |    ```bash
223 |    git push origin feature/your-feature-name
224 |    ```
225 | 
226 | 5. **Create Pull Request** on GitHub
227 | 
228 | ### PR Guidelines
229 | 
230 | - Write clear, descriptive commit messages
231 | - Include tests for new features
232 | - Update documentation as needed
233 | - Keep PRs focused on a single feature/fix
234 | - Respond to review feedback promptly
235 | 
236 | ## Code Style
237 | 
238 | - Follow PEP 8
239 | - Use type hints for all functions
240 | - Write docstrings in Google style
241 | - Keep functions small and focused
242 | - Use meaningful variable names
243 | 
244 | ## Questions?
245 | 
246 | - Open an issue on GitHub
247 | - Join discussions in pull requests
248 | - Check existing documentation
249 | 
250 | Thank you for contributing! 🎉
251 | 


--------------------------------------------------------------------------------
/querygym/methods/genqr_ensemble.py:
--------------------------------------------------------------------------------
  1 | from __future__ import annotations
  2 | from ..core.base import BaseReformulator, QueryItem, ReformulationResult
  3 | from ..core.registry import register_method
  4 | from typing import List
  5 | 
  6 | # 10 instruction variants from GenQREnsemble paper
  7 | VARIANT_IDS = [
  8 |     "genqr_ensemble.inst1.v1",
  9 |     "genqr_ensemble.inst2.v1",
 10 |     "genqr_ensemble.inst3.v1",
 11 |     "genqr_ensemble.inst4.v1",
 12 |     "genqr_ensemble.inst5.v1",
 13 |     "genqr_ensemble.inst6.v1",
 14 |     "genqr_ensemble.inst7.v1",
 15 |     "genqr_ensemble.inst8.v1",
 16 |     "genqr_ensemble.inst9.v1",
 17 |     "genqr_ensemble.inst10.v1",
 18 | ]
 19 | 
 20 | @register_method("genqr_ensemble")
 21 | class GenQREnsemble(BaseReformulator):
 22 |     """
 23 |     GenQREnsemble: Generative Query Reformulation with Ensemble of Instructions.
 24 |     
 25 |     Uses 10 instruction variants to generate diverse keyword expansions.
 26 |     Each variant generates keywords independently, then all are merged.
 27 |     
 28 |     Pipeline:
 29 |         1. For each of 10 instruction variants, generate keywords (10 LLM calls)
 30 |         2. Parse and merge all keyword lists
 31 |         3. Expand query: (Q × 5) + K1 + K2 + ... + K10
 32 |     
 33 |     Sampling Parameters (from paper):
 34 |         - temperature: 0.92 (nucleus sampling)
 35 |         - max_tokens: 256
 36 |         - parallel: false (can be enabled for concurrent generation)
 37 |     
 38 |     Config Parameters:
 39 |         - variant_ids: List of prompt IDs (defaults to all 10 variants)
 40 |         - parallel: Enable parallel generation of all variants (default: false)
 41 |     """
 42 |     VERSION = "1.0"
 43 |     CONCATENATION_STRATEGY = "query_repeat_plus_generated"  # (Q × 5) + K1 + K2 + ... + K10
 44 |     DEFAULT_QUERY_REPEATS = 5  # GenQREnsemble uses 5 repetitions
 45 | 
 46 |     def _parse_keywords(self, keywords_text: str) -> List[str]:
 47 |         """
 48 |         Parse keywords from LLM output (comma-separated or bullet list).
 49 |         
 50 |         Args:
 51 |             keywords_text: Raw LLM output
 52 |             
 53 |         Returns:
 54 |             List of individual keywords
 55 |         """
 56 |         keywords = []
 57 |         
 58 |         # Check if it's a bullet list (contains dashes/newlines)
 59 |         if '\n' in keywords_text or keywords_text.strip().startswith('-'):
 60 |             # Split by newlines and extract after dashes
 61 |             for line in keywords_text.split('\n'):
 62 |                 line = line.strip()
 63 |                 # Remove leading dash/bullet
 64 |                 if line.startswith('-'):
 65 |                     line = line[1:].strip()
 66 |                 elif line.startswith('•'):
 67 |                     line = line[1:].strip()
 68 |                 elif line.startswith('*'):
 69 |                     line = line[1:].strip()
 70 |                 
 71 |                 if line:
 72 |                     keywords.append(line)
 73 |         else:
 74 |             # Split by comma (standard format)
 75 |             keywords = [k.strip() for k in keywords_text.split(',') if k.strip()]
 76 |         
 77 |         return keywords
 78 | 
 79 |     def reformulate(self, q: QueryItem, contexts=None) -> ReformulationResult:
 80 |         """Generate keywords using 10 instruction variants and merge."""
 81 |         # Get parameters (matching baseline sampling params)
 82 |         temperature = float(self.cfg.llm.get("temperature", 0.92))
 83 |         max_tokens = int(self.cfg.llm.get("max_tokens", 256))
 84 |         parallel = bool(self.cfg.params.get("parallel", False))
 85 |         
 86 |         # Get variant IDs (configurable, defaults to all 10)
 87 |         variant_ids = self.cfg.params.get("variant_ids", VARIANT_IDS)
 88 |         
 89 |         # Generate keywords for each instruction variant
 90 |         all_keywords = []
 91 |         variant_outputs = {}
 92 |         
 93 |         if parallel:
 94 |             # Parallel generation using ThreadPoolExecutor
 95 |             from concurrent.futures import ThreadPoolExecutor, as_completed
 96 |             
 97 |             def generate_keywords_for_variant(idx: int, prompt_id: str):
 98 |                 msgs = self.prompts.render(prompt_id, query=q.text)
 99 |                 output = self.llm.chat(msgs, temperature=temperature, max_tokens=max_tokens)
100 |                 keywords = self._parse_keywords(output)
101 |                 return idx, prompt_id, output, keywords
102 |             
103 |             # Generate all variants in parallel
104 |             with ThreadPoolExecutor(max_workers=len(variant_ids)) as executor:
105 |                 futures = [
106 |                     executor.submit(generate_keywords_for_variant, i, prompt_id)
107 |                     for i, prompt_id in enumerate(variant_ids, 1)
108 |                 ]
109 |                 for future in as_completed(futures):
110 |                     i, prompt_id, output, keywords = future.result()
111 |                     all_keywords.extend(keywords)
112 |                     variant_outputs[f"variant_{i}"] = {
113 |                         "prompt_id": prompt_id,
114 |                         "raw_output": output,
115 |                         "keywords": keywords,
116 |                         "count": len(keywords)
117 |                     }
118 |         else:
119 |             # Sequential generation (default)
120 |             for i, prompt_id in enumerate(variant_ids, 1):
121 |                 # Render prompt and generate keywords
122 |                 msgs = self.prompts.render(prompt_id, query=q.text)
123 |                 output = self.llm.chat(msgs, temperature=temperature, max_tokens=max_tokens)
124 |                 
125 |                 # Parse keywords from output
126 |                 keywords = self._parse_keywords(output)
127 |                 all_keywords.extend(keywords)
128 |                 
129 |                 # Store per-variant output for metadata
130 |                 variant_outputs[f"variant_{i}"] = {
131 |                     "prompt_id": prompt_id,
132 |                     "raw_output": output,
133 |                     "keywords": keywords,
134 |                     "count": len(keywords)
135 |                 }
136 |         
137 |         # Merge all keywords into single string
138 |         generated_content = " ".join(all_keywords)
139 |         
140 |         # Use base class concatenation: (Q × 5) + generated_content
141 |         reformulated = self.concatenate_result(q.text, generated_content)
142 |         
143 |         return ReformulationResult(
144 |             q.qid, 
145 |             q.text, 
146 |             reformulated,
147 |             metadata={
148 |                 "num_variants": len(variant_ids),
149 |                 "total_keywords": len(all_keywords),
150 |                 "keywords": all_keywords,
151 |                 "variant_outputs": variant_outputs,
152 |                 "temperature": temperature,
153 |                 "max_tokens": max_tokens,
154 |                 "parallel": parallel
155 |             }
156 |         )
157 | 


--------------------------------------------------------------------------------
/querygym/adapters/pyterrier_adapter.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | """
  3 | PyTerrier Searcher Adapter for querygym
  4 | 
  5 | This module provides an adapter that wraps PyTerrier's search functionality
  6 | to implement the querygym BaseSearcher interface.
  7 | """
  8 | 
  9 | import pandas as pd
 10 | from typing import List, Dict, Any, Optional
 11 | from ..core.searcher import BaseSearcher, SearchHit, SearcherRegistry
 12 | 
 13 | 
 14 | class PyTerrierSearcher(BaseSearcher):
 15 |     """
 16 |     PyTerrier adapter implementing the BaseSearcher interface.
 17 |     
 18 |     This adapter wraps PyTerrier's search functionality to provide
 19 |     a standardized interface for querygym.
 20 |     """
 21 |     
 22 |     def __init__(self, index_path: Optional[str] = None, 
 23 |                  index_name: Optional[str] = None,
 24 |                  searcher_type: str = "bm25", **kwargs):
 25 |         """
 26 |         Initialize PyTerrier searcher.
 27 |         
 28 |         Args:
 29 |             index_path: Path to PyTerrier index
 30 |             index_name: Name of prebuilt PyTerrier index
 31 |             searcher_type: Type of searcher ("bm25", "tfidf", "pl2", etc.)
 32 |             **kwargs: Additional arguments
 33 |         """
 34 |         try:
 35 |             import pyterrier as pt
 36 |         except ImportError:
 37 |             raise ImportError("PyTerrier is required for PyTerrierSearcher. Install with: pip install python-terrier")
 38 |         
 39 |         self.searcher_type = searcher_type
 40 |         self._searcher_info = {
 41 |             "name": "PyTerrierSearcher",
 42 |             "type": searcher_type,
 43 |             "index_path": index_path,
 44 |             "index_name": index_name
 45 |         }
 46 |         
 47 |         # Initialize PyTerrier if not already done
 48 |         if not pt.started():
 49 |             pt.init()
 50 |         
 51 |         # Load index
 52 |         if index_path:
 53 |             self.index = pt.IndexRef.of(index_path)
 54 |         elif index_name:
 55 |             # Try to load prebuilt index
 56 |             try:
 57 |                 self.index = pt.get_dataset(index_name).get_index()
 58 |             except Exception as e:
 59 |                 raise ValueError(f"Could not load prebuilt index '{index_name}': {e}")
 60 |         else:
 61 |             raise ValueError("Either index_path or index_name must be provided")
 62 |         
 63 |         # Create searcher based on type
 64 |         self.searcher = self._create_searcher(searcher_type, **kwargs)
 65 |     
 66 |     def _create_searcher(self, searcher_type: str, **kwargs):
 67 |         """Create the appropriate searcher based on type."""
 68 |         import pyterrier as pt
 69 |         
 70 |         # Create a proper PyTerrier Retriever, not a pipeline component
 71 |         if searcher_type == "bm25":
 72 |             return pt.terrier.Retriever(self.index, wmodel="BM25", **kwargs)
 73 |         elif searcher_type == "tfidf":
 74 |             return pt.terrier.Retriever(self.index, wmodel="TF_IDF", **kwargs)
 75 |         elif searcher_type == "pl2":
 76 |             return pt.terrier.Retriever(self.index, wmodel="PL2", **kwargs)
 77 |         elif searcher_type == "dph":
 78 |             return pt.terrier.Retriever(self.index, wmodel="DPH", **kwargs)
 79 |         elif searcher_type == "dirichletlm":
 80 |             return pt.terrier.Retriever(self.index, wmodel="DirichletLM", **kwargs)
 81 |         else:
 82 |             # Default to BM25
 83 |             return pt.terrier.Retriever(self.index, wmodel="BM25", **kwargs)
 84 |     
 85 |     def search(self, query: str, k: int = 10, **kwargs) -> List[SearchHit]:
 86 |         """Search for documents using a single query."""
 87 |         # Use PyTerrier's search method which returns a DataFrame
 88 |         search_results = self.searcher.search(query)
 89 |         
 90 |         # Limit to top-k results
 91 |         search_results = search_results.head(k)
 92 |         
 93 |         return self._process_results(search_results)
 94 |     
 95 |     def batch_search(self, queries: List[str], k: int = 10, 
 96 |                     num_threads: int = 1, **kwargs) -> List[List[SearchHit]]:
 97 |         """Search for documents using multiple queries in batch."""
 98 |         # Create query DataFrame for batch processing
 99 |         query_df = pd.DataFrame({
100 |             'qid': [str(i) for i in range(len(queries))],
101 |             'query': queries
102 |         })
103 |         
104 |         # Use PyTerrier's transform method for batch processing
105 |         search_results = self.searcher.transform(query_df)
106 |         
107 |         # Group results by query ID
108 |         batch_results = []
109 |         for qid in range(len(queries)):
110 |             query_results = search_results[search_results['qid'] == str(qid)].head(k)
111 |             batch_results.append(self._process_results(query_results))
112 |         
113 |         return batch_results
114 |     
115 |     def _process_results(self, results_df: pd.DataFrame) -> List[SearchHit]:
116 |         """Process PyTerrier results into SearchHit objects."""
117 |         search_hits = []
118 |         
119 |         if results_df.empty:
120 |             return search_hits
121 |         
122 |         for _, row in results_df.iterrows():
123 |             try:
124 |                 # Handle different possible column names
125 |                 docid = str(row.get('docno', row.get('docid', '')))
126 |                 score = float(row.get('score', 0.0))
127 |                 
128 |                 # Try different content field names
129 |                 content = (row.get('text', '') or 
130 |                           row.get('content', '') or 
131 |                           row.get('body', '') or 
132 |                           str(docid))
133 |                 
134 |                 search_hit = SearchHit(
135 |                     docid=docid,
136 |                     score=score,
137 |                     content=str(content),
138 |                     metadata={
139 |                         'qid': str(row.get('qid', '')),
140 |                         'docno': str(row.get('docno', '')),
141 |                         'searcher_type': self.searcher_type,
142 |                         'available_columns': list(row.index)
143 |                     }
144 |                 )
145 |                 search_hits.append(search_hit)
146 |                 
147 |             except Exception as e:
148 |                 # Fallback for malformed data
149 |                 search_hit = SearchHit(
150 |                     docid=str(row.get('docno', 'unknown')),
151 |                     score=0.0,
152 |                     content=str(row.get('docno', 'unknown')),
153 |                     metadata={
154 |                         'error': str(e),
155 |                         'fallback': True,
156 |                         'searcher_type': self.searcher_type
157 |                     }
158 |                 )
159 |                 search_hits.append(search_hit)
160 |         
161 |         return search_hits
162 |     
163 |     def get_searcher_info(self) -> Dict[str, Any]:
164 |         """Get information about the searcher configuration."""
165 |         return self._searcher_info.copy()
166 |     
167 |     def configure(self, **kwargs) -> None:
168 |         """Configure searcher parameters."""
169 |         # Update searcher info
170 |         self._searcher_info.update(kwargs)
171 |         
172 |         # Recreate searcher with new parameters
173 |         self.searcher = self._create_searcher(self.searcher_type, **kwargs)
174 | 
175 | 
176 | # Register the searcher
177 | SearcherRegistry.register("pyterrier", PyTerrierSearcher)
178 | 


--------------------------------------------------------------------------------
/querygym/__init__.py:
--------------------------------------------------------------------------------
  1 | """
  2 | querygym: LLM-based Query Reformulation Toolkit
  3 | 
  4 | Simple usage:
  5 |     import querygym as qg
  6 |     
  7 |     # Create a reformulator
  8 |     reformulator = qg.create_reformulator("genqr_ensemble")
  9 |     
 10 |     # Reformulate queries
 11 |     result = reformulator.reformulate(qg.QueryItem("q1", "what causes diabetes"))
 12 | """
 13 | 
 14 | __version__ = "0.1.0"
 15 | 
 16 | # Core data structures
 17 | from .core.base import QueryItem, ReformulationResult, MethodConfig, BaseReformulator
 18 | 
 19 | # LLM and prompts
 20 | from .core.llm import OpenAICompatibleClient
 21 | from .core.prompts import PromptBank
 22 | 
 23 | # Searcher interface
 24 | from .core.searcher import BaseSearcher, SearchHit, SearcherRegistry, create_searcher
 25 | 
 26 | # Searcher wrappers for user convenience
 27 | from .core.searcher_wrappers import wrap_pyserini_searcher, wrap_pyterrier_retriever, wrap_custom_searcher
 28 | 
 29 | # Data loaders
 30 | from .data.dataloader import DataLoader, UnifiedQuerySource
 31 | 
 32 | # Format-specific loaders
 33 | from . import loaders
 34 | 
 35 | # All reformulation methods
 36 | from .methods import (
 37 |     GENQR,
 38 |     GenQREnsemble,
 39 |     Query2Doc,
 40 |     QAExpand,
 41 |     MuGI,
 42 |     LameR,
 43 |     Query2E,
 44 |     CSQE,
 45 | )
 46 | 
 47 | # High-level runner
 48 | from .core.runner import run_method, build_llm
 49 | from .core.registry import METHODS, register_method
 50 | 
 51 | 
 52 | # Import adapters to register them
 53 | from . import adapters
 54 | 
 55 | # Convenience factory function
 56 | def create_reformulator(
 57 |     method_name: str,
 58 |     model: str = "gpt-4",
 59 |     params: dict = None,
 60 |     llm_config: dict = None,
 61 |     prompt_bank_path: str = None,
 62 |     **kwargs
 63 | ):
 64 |     """
 65 |     Create a reformulator instance with sensible defaults.
 66 |     
 67 |     Args:
 68 |         method_name: Name of the method (e.g., "genqr", "genqr_ensemble", "query2doc")
 69 |         model: LLM model name (default: "gpt-4")
 70 |         params: Method-specific parameters (default: {})
 71 |         llm_config: Additional LLM configuration (temperature, max_tokens, etc.)
 72 |         prompt_bank_path: Path to prompt bank YAML (default: bundled prompt_bank.yaml)
 73 |         **kwargs: Additional MethodConfig parameters (seed, retries)
 74 |     
 75 |     Returns:
 76 |         BaseReformulator instance
 77 |     
 78 |     Example:
 79 |         >>> import querygym as qg
 80 |         >>> reformulator = qg.create_reformulator("genqr_ensemble", model="gpt-4")
 81 |         >>> result = reformulator.reformulate(qg.QueryItem("q1", "what causes diabetes"))
 82 |     """
 83 |     from pathlib import Path
 84 |     
 85 |     if params is None:
 86 |         params = {}
 87 |     
 88 |     if llm_config is None:
 89 |         llm_config = {}
 90 |     
 91 |     # Set up LLM config
 92 |     llm_cfg = {
 93 |         "model": model,
 94 |         "temperature": llm_config.get("temperature", 0.8),
 95 |         "max_tokens": llm_config.get("max_tokens", 256),
 96 |         **{k: v for k, v in llm_config.items() if k not in ["temperature", "max_tokens"]}
 97 |     }
 98 |     
 99 |     # Create config
100 |     config = MethodConfig(
101 |         name=method_name,
102 |         params=params,
103 |         llm=llm_cfg,
104 |         seed=kwargs.get("seed", 42),
105 |         retries=kwargs.get("retries", 2)
106 |     )
107 |     
108 |     # Build LLM client
109 |     llm = build_llm(config)
110 |     
111 |     # Load prompt bank
112 |     if prompt_bank_path is None:
113 |         prompt_bank_path = Path(__file__).parent / "prompt_bank.yaml"
114 |     pb = PromptBank(prompt_bank_path)
115 |     
116 |     # Get method class and instantiate
117 |     if method_name not in METHODS:
118 |         raise ValueError(f"Unknown method: {method_name}. Available: {list(METHODS.keys())}")
119 |     
120 |     MethodClass = METHODS[method_name]
121 |     return MethodClass(config, llm, pb)
122 | 
123 | 
124 | def load_queries(path: str, format: str = "tsv", **kwargs):
125 |     """
126 |     Load queries from a local file.
127 |     
128 |     Args:
129 |         path: Path to queries file
130 |         format: File format - "tsv" or "jsonl" (default: "tsv")
131 |         **kwargs: Additional parameters for DataLoader.load_queries()
132 |     
133 |     Returns:
134 |         List of QueryItem objects
135 |     
136 |     Example:
137 |         >>> queries = qg.load_queries("queries.tsv", format="tsv")
138 |         >>> queries = qg.load_queries("queries.jsonl", format="jsonl")
139 |     """
140 |     return DataLoader.load_queries(path, format=format, **kwargs)
141 | 
142 | 
143 | def load_qrels(path: str, format: str = "trec", **kwargs):
144 |     """
145 |     Load qrels from a local file.
146 |     
147 |     Args:
148 |         path: Path to qrels file
149 |         format: File format - "trec" (default: "trec")
150 |         **kwargs: Additional parameters for DataLoader.load_qrels()
151 |     
152 |     Returns:
153 |         Dict mapping qid -> {docid -> relevance}
154 |     
155 |     Example:
156 |         >>> qrels = qg.load_qrels("qrels.txt")
157 |     """
158 |     return DataLoader.load_qrels(path, format=format, **kwargs)
159 | 
160 | 
161 | def load_contexts(path: str, **kwargs):
162 |     """
163 |     Load contexts from a JSONL file.
164 |     
165 |     Args:
166 |         path: Path to contexts JSONL file
167 |         **kwargs: Additional parameters for DataLoader.load_contexts()
168 |     
169 |     Returns:
170 |         Dict mapping qid -> list of context strings
171 |     
172 |     Example:
173 |         >>> contexts = qg.load_contexts("contexts.jsonl")
174 |     """
175 |     return DataLoader.load_contexts(path, **kwargs)
176 | 
177 | 
178 | def load_examples(path: str, **kwargs):
179 |     """
180 |     Load few-shot examples from a JSONL file.
181 |     
182 |     Used for methods like Query2E that need (query, passage) pairs as demonstrations.
183 |     
184 |     Args:
185 |         path: Path to examples JSONL file
186 |         **kwargs: Additional parameters for DataLoader.load_examples()
187 |     
188 |     Returns:
189 |         List of dicts with 'query' and 'passage' keys
190 |     
191 |     Example:
192 |         >>> examples = qg.load_examples("examples.jsonl")
193 |         >>> reformulator = qg.create_reformulator("query2e", params={"mode": "fs", "examples": examples})
194 |         
195 |     JSONL format:
196 |         {"query": "how long is flea life cycle?", "passage": "The life cycle of a flea..."}
197 |         {"query": "cost of flooring?", "passage": "The cost of interior concrete..."}
198 |     """
199 |     return DataLoader.load_examples(path, **kwargs)
200 | 
201 | 
202 | __all__ = [
203 |     # Version
204 |     "__version__",
205 |     
206 |     # Core classes
207 |     "QueryItem",
208 |     "ReformulationResult",
209 |     "MethodConfig",
210 |     "BaseReformulator",
211 |     
212 |     # LLM & Prompts
213 |     "OpenAICompatibleClient",
214 |     "PromptBank",
215 |     
216 |     # Searcher interface
217 |     "BaseSearcher",
218 |     "SearchHit",
219 |     "SearcherRegistry",
220 |     "create_searcher",
221 |     
222 |     # Searcher wrappers
223 |     "wrap_pyserini_searcher",
224 |     "wrap_pyterrier_retriever", 
225 |     "wrap_custom_searcher",
226 |     
227 |     # Data
228 |     "DataLoader",
229 |     "UnifiedQuerySource",  # Deprecated
230 |     "loaders",
231 |     
232 |     # Methods
233 |     "GENQR",
234 |     "GenQREnsemble",
235 |     "Query2Doc",
236 |     "QAExpand",
237 |     "MuGI",
238 |     "LameR",
239 |     "Query2E",
240 |     "CSQE",
241 |     
242 |     # High-level API
243 |     "run_method",
244 |     "build_llm",
245 |     "create_reformulator",
246 |     "load_queries",
247 |     "load_qrels",
248 |     "load_contexts",
249 |     "load_examples",
250 |     
251 |     # Retrieval
252 |     "Retriever",
253 |     
254 |     # Registry
255 |     "METHODS",
256 |     "register_method",
257 | ]
258 | 


--------------------------------------------------------------------------------
/docs/api/loaders.md:
--------------------------------------------------------------------------------
  1 | # Data Loaders API Reference
  2 | 
  3 | ## Overview
  4 | 
  5 | querygym provides loaders for common IR dataset formats. All loaders return standardized `QueryItem` objects and dictionaries for easy integration.
  6 | 
  7 | ## Core DataLoader
  8 | 
  9 | The `DataLoader` class provides format-agnostic loading utilities.
 10 | 
 11 | ### load_queries
 12 | 
 13 | Load queries from TSV or JSONL files.
 14 | 
 15 | ```python
 16 | from querygym.data.dataloader import DataLoader
 17 | 
 18 | # Load from TSV (qid<TAB>text format)
 19 | queries = DataLoader.load_queries("queries.tsv")
 20 | 
 21 | # Load from JSONL with custom keys
 22 | queries = DataLoader.load_queries(
 23 |     "queries.jsonl",
 24 |     format="jsonl",
 25 |     qid_key="id",
 26 |     query_key="text"
 27 | )
 28 | ```
 29 | 
 30 | **Parameters:**
 31 | - `path` (str): Path to queries file
 32 | - `format` (str): File format - "tsv" or "jsonl" (default: "tsv")
 33 | - `qid_col` (int): Column index for query ID in TSV (default: 0)
 34 | - `query_col` (int): Column index for query text in TSV (default: 1)
 35 | - `qid_key` (str): JSON key for query ID in JSONL (default: "qid")
 36 | - `query_key` (str): JSON key for query text in JSONL (default: "text")
 37 | 
 38 | **Returns:** `List[QueryItem]`
 39 | 
 40 | ### load_qrels
 41 | 
 42 | Load relevance judgments (qrels).
 43 | 
 44 | ```python
 45 | # Load TREC format qrels (qid iter docid relevance)
 46 | qrels = DataLoader.load_qrels("qrels.txt", format="trec")
 47 | ```
 48 | 
 49 | **Parameters:**
 50 | - `path` (str): Path to qrels file
 51 | - `format` (str): File format - "trec" (default: "trec")
 52 | 
 53 | **Returns:** `Dict[str, Dict[str, int]]` - Nested dict: qid → docid → relevance
 54 | 
 55 | ### load_contexts
 56 | 
 57 | Load retrieved contexts for query reformulation.
 58 | 
 59 | ```python
 60 | # Load contexts from JSONL
 61 | contexts = DataLoader.load_contexts("contexts.jsonl")
 62 | ```
 63 | 
 64 | **Parameters:**
 65 | - `path` (str): Path to contexts JSONL file
 66 | 
 67 | **Returns:** `Dict[str, List[str]]` - Dict: qid → list of context strings
 68 | 
 69 | ### save_queries
 70 | 
 71 | Save queries to file.
 72 | 
 73 | ```python
 74 | # Save to TSV
 75 | DataLoader.save_queries(queries, "output.tsv")
 76 | 
 77 | # Save to JSONL
 78 | DataLoader.save_queries(queries, "output.jsonl", format="jsonl")
 79 | ```
 80 | 
 81 | **Parameters:**
 82 | - `queries` (List[QueryItem]): Queries to save
 83 | - `path` (str): Output file path
 84 | - `format` (str): Output format - "tsv" or "jsonl" (default: "tsv")
 85 | 
 86 | ---
 87 | 
 88 | ## BEIR Format Loaders
 89 | 
 90 | Loaders for [BEIR](https://github.com/beir-cellar/beir) benchmark datasets.
 91 | 
 92 | ### Installation
 93 | 
 94 | ```bash
 95 | pip install querygym[beir]
 96 | ```
 97 | 
 98 | ### Download BEIR Datasets
 99 | 
100 | ```python
101 | from beir.datasets.data_loader import GenericDataLoader
102 | 
103 | # Download and extract a BEIR dataset
104 | data_path = GenericDataLoader("nfcorpus").download_and_unzip()
105 | ```
106 | 
107 | ### querygym.loaders.beir.load_queries
108 | 
109 | Load queries from a BEIR dataset directory.
110 | 
111 | ```python
112 | from querygym.loaders.beir import load_queries
113 | 
114 | # Load queries (uses queries.jsonl)
115 | queries = load_queries("./data/nfcorpus")
116 | ```
117 | 
118 | **Parameters:**
119 | - `beir_data_dir` (str | Path): Path to BEIR dataset directory
120 | - `split` (str): Dataset split (default: "test", kept for API consistency)
121 | 
122 | **Returns:** `List[QueryItem]`
123 | 
124 | **File Format:** BEIR uses `queries.jsonl` with `{"_id": "...", "text": "..."}`
125 | 
126 | ### querygym.loaders.beir.load_qrels
127 | 
128 | Load qrels from a BEIR dataset.
129 | 
130 | ```python
131 | from querygym.loaders.beir import load_qrels
132 | 
133 | # Load test split qrels
134 | qrels = load_qrels("./data/nfcorpus", split="test")
135 | 
136 | # Load dev split
137 | qrels = load_qrels("./data/nfcorpus", split="dev")
138 | ```
139 | 
140 | **Parameters:**
141 | - `beir_data_dir` (str | Path): Path to BEIR dataset directory
142 | - `split` (str): Dataset split - "train", "dev", or "test" (default: "test")
143 | 
144 | **Returns:** `Dict[str, Dict[str, int]]`
145 | 
146 | **File Location:** `{beir_data_dir}/qrels/{split}.tsv`
147 | 
148 | ### querygym.loaders.beir.load_corpus
149 | 
150 | Load document corpus from a BEIR dataset.
151 | 
152 | ```python
153 | from querygym.loaders.beir import load_corpus
154 | 
155 | # Load corpus
156 | corpus = load_corpus("./data/nfcorpus")
157 | 
158 | # Access documents
159 | doc = corpus["doc_id_123"]
160 | print(doc["title"], doc["text"])
161 | ```
162 | 
163 | **Parameters:**
164 | - `beir_data_dir` (str | Path): Path to BEIR dataset directory
165 | 
166 | **Returns:** `Dict[str, Dict[str, str]]` - Dict: doc_id → {"title": ..., "text": ...}
167 | 
168 | **File Format:** BEIR uses `corpus.jsonl` with `{"_id": "...", "title": "...", "text": "..."}`
169 | 
170 | ---
171 | 
172 | ## MS MARCO Format Loaders
173 | 
174 | Loaders for [MS MARCO](https://microsoft.github.io/msmarco/) datasets.
175 | 
176 | ### Installation
177 | 
178 | ```bash
179 | pip install querygym  # No extra dependencies needed
180 | ```
181 | 
182 | ### Download MS MARCO Datasets
183 | 
184 | ```python
185 | import ir_datasets
186 | 
187 | # Download via ir_datasets
188 | dataset = ir_datasets.load("msmarco-passage/dev")
189 | 
190 | # Or download from official sources
191 | # https://microsoft.github.io/msmarco/
192 | ```
193 | 
194 | ### querygym.loaders.msmarco.load_queries
195 | 
196 | Load queries from MS MARCO TSV file.
197 | 
198 | ```python
199 | from querygym.loaders.msmarco import load_queries
200 | 
201 | # Load queries
202 | queries = load_queries("./data/queries.dev.tsv")
203 | ```
204 | 
205 | **Parameters:**
206 | - `queries_tsv` (str | Path): Path to MS MARCO queries TSV file
207 | 
208 | **Returns:** `List[QueryItem]`
209 | 
210 | **File Format:** TSV with `qid<TAB>query_text`
211 | 
212 | ### querygym.loaders.msmarco.load_qrels
213 | 
214 | Load qrels from MS MARCO file.
215 | 
216 | ```python
217 | from querygym.loaders.msmarco import load_qrels
218 | 
219 | # Load qrels (handles both TREC and simplified formats)
220 | qrels = load_qrels("./data/qrels.dev.tsv")
221 | ```
222 | 
223 | **Parameters:**
224 | - `qrels_tsv` (str | Path): Path to MS MARCO qrels file
225 | 
226 | **Returns:** `Dict[str, Dict[str, int]]`
227 | 
228 | **File Formats Supported:**
229 | - TREC format: `qid<TAB>0<TAB>docid<TAB>relevance`
230 | - Simplified: `qid<TAB>docid<TAB>relevance`
231 | 
232 | ### querygym.loaders.msmarco.load_collection
233 | 
234 | Load document collection from MS MARCO TSV file.
235 | 
236 | ```python
237 | from querygym.loaders.msmarco import load_collection
238 | 
239 | # Load collection
240 | collection = load_collection("./data/collection.tsv")
241 | 
242 | # Access documents
243 | doc_text = collection["doc_id_123"]
244 | ```
245 | 
246 | **Parameters:**
247 | - `collection_tsv` (str | Path): Path to MS MARCO collection TSV file
248 | 
249 | **Returns:** `Dict[str, str]` - Dict: doc_id → document_text
250 | 
251 | **File Format:** TSV with `docid<TAB>text`
252 | 
253 | ---
254 | 
255 | ## Complete Example
256 | 
257 | ```python
258 | import querygym as qg
259 | 
260 | # Load BEIR dataset
261 | queries = qg.loaders.beir.load_queries("./data/nfcorpus")
262 | qrels = qg.loaders.beir.load_qrels("./data/nfcorpus", split="test")
263 | corpus = qg.loaders.beir.load_corpus("./data/nfcorpus")
264 | 
265 | # Or load MS MARCO dataset
266 | queries = qg.loaders.msmarco.load_queries("./data/queries.tsv")
267 | qrels = qg.loaders.msmarco.load_qrels("./data/qrels.tsv")
268 | collection = qg.loaders.msmarco.load_collection("./data/collection.tsv")
269 | 
270 | # Use with reformulation
271 | reformulator = qg.create_reformulator("genqr_ensemble", model="gpt-4")
272 | results = reformulator.reformulate_batch(queries)
273 | 
274 | # Save reformulated queries
275 | qg.DataLoader.save_queries(
276 |     [qg.QueryItem(r.qid, r.reformulated) for r in results],
277 |     "reformulated.tsv"
278 | )
279 | ```
280 | 


--------------------------------------------------------------------------------
/examples/docker/quickstart.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |  "cells": [
  3 |   {
  4 |    "cell_type": "markdown",
  5 |    "metadata": {},
  6 |    "source": [
  7 |     "# QueryGym Docker Quick Start\n",
  8 |     "\n",
  9 |     "This notebook demonstrates basic QueryGym usage within the Docker container.\n",
 10 |     "\n",
 11 |     "## Prerequisites\n",
 12 |     "- Container is running with mounted volumes\n",
 13 |     "- OPENAI_API_KEY environment variable is set"
 14 |    ]
 15 |   },
 16 |   {
 17 |    "cell_type": "code",
 18 |    "execution_count": null,
 19 |    "metadata": {},
 20 |    "outputs": [],
 21 |    "source": [
 22 |     "# Import QueryGym\n",
 23 |     "import querygym as qg\n",
 24 |     "import os\n",
 25 |     "\n",
 26 |     "print(f\"QueryGym version: {qg.__version__}\")\n",
 27 |     "print(f\"API key configured: {'Yes' if os.getenv('OPENAI_API_KEY') else 'No'}\")"
 28 |    ]
 29 |   },
 30 |   {
 31 |    "cell_type": "markdown",
 32 |    "metadata": {},
 33 |    "source": [
 34 |     "## 1. Create Sample Queries"
 35 |    ]
 36 |   },
 37 |   {
 38 |    "cell_type": "code",
 39 |    "execution_count": null,
 40 |    "metadata": {},
 41 |    "outputs": [],
 42 |    "source": [
 43 |     "# Create sample queries\n",
 44 |     "sample_queries = [\n",
 45 |     "    qg.QueryItem(qid=\"1\", text=\"what is covid-19\"),\n",
 46 |     "    qg.QueryItem(qid=\"2\", text=\"how does machine learning work\"),\n",
 47 |     "    qg.QueryItem(qid=\"3\", text=\"best practices for docker containers\")\n",
 48 |     "]\n",
 49 |     "\n",
 50 |     "print(f\"Created {len(sample_queries)} sample queries\")\n",
 51 |     "for q in sample_queries:\n",
 52 |     "    print(f\"  {q.qid}: {q.text}\")"
 53 |    ]
 54 |   },
 55 |   {
 56 |    "cell_type": "markdown",
 57 |    "metadata": {},
 58 |    "source": [
 59 |     "## 2. Create a Reformulator"
 60 |    ]
 61 |   },
 62 |   {
 63 |    "cell_type": "code",
 64 |    "execution_count": null,
 65 |    "metadata": {},
 66 |    "outputs": [],
 67 |    "source": [
 68 |     "# Create a reformulator using GenQR method\n",
 69 |     "reformulator = qg.create_reformulator(\n",
 70 |     "    method=\"genqr\",\n",
 71 |     "    model=\"gpt-3.5-turbo\",  # or \"gpt-4\" for better quality\n",
 72 |     "    temperature=0.7\n",
 73 |     ")\n",
 74 |     "\n",
 75 |     "print(f\"Created reformulator: {reformulator.__class__.__name__}\")"
 76 |    ]
 77 |   },
 78 |   {
 79 |    "cell_type": "markdown",
 80 |    "metadata": {},
 81 |    "source": [
 82 |     "## 3. Reformulate a Single Query"
 83 |    ]
 84 |   },
 85 |   {
 86 |    "cell_type": "code",
 87 |    "execution_count": null,
 88 |    "metadata": {},
 89 |    "outputs": [],
 90 |    "source": [
 91 |     "# Reformulate the first query\n",
 92 |     "result = reformulator.reformulate(sample_queries[0])\n",
 93 |     "\n",
 94 |     "print(\"Original query:\", result.original)\n",
 95 |     "print(\"Reformulated query:\", result.reformulated)\n",
 96 |     "print(\"Metadata:\", result.metadata)"
 97 |    ]
 98 |   },
 99 |   {
100 |    "cell_type": "markdown",
101 |    "metadata": {},
102 |    "source": [
103 |     "## 4. Batch Reformulation"
104 |    ]
105 |   },
106 |   {
107 |    "cell_type": "code",
108 |    "execution_count": null,
109 |    "metadata": {},
110 |    "outputs": [],
111 |    "source": [
112 |     "# Reformulate all queries\n",
113 |     "results = reformulator.reformulate_batch(sample_queries)\n",
114 |     "\n",
115 |     "print(f\"Reformulated {len(results)} queries:\\n\")\n",
116 |     "for r in results:\n",
117 |     "    print(f\"QID {r.qid}:\")\n",
118 |     "    print(f\"  Original: {r.original}\")\n",
119 |     "    print(f\"  Reformulated: {r.reformulated}\")\n",
120 |     "    print()"
121 |    ]
122 |   },
123 |   {
124 |    "cell_type": "markdown",
125 |    "metadata": {},
126 |    "source": [
127 |     "## 5. Save Results"
128 |    ]
129 |   },
130 |   {
131 |    "cell_type": "code",
132 |    "execution_count": null,
133 |    "metadata": {},
134 |    "outputs": [],
135 |    "source": [
136 |     "# Save reformulated queries to outputs directory\n",
137 |     "output_queries = [qg.QueryItem(r.qid, r.reformulated) for r in results]\n",
138 |     "qg.DataLoader.save_queries(output_queries, \"outputs/reformulated_queries.tsv\")\n",
139 |     "\n",
140 |     "print(\"Saved reformulated queries to: outputs/reformulated_queries.tsv\")"
141 |    ]
142 |   },
143 |   {
144 |    "cell_type": "markdown",
145 |    "metadata": {},
146 |    "source": [
147 |     "## 6. Explore Available Methods"
148 |    ]
149 |   },
150 |   {
151 |    "cell_type": "code",
152 |    "execution_count": null,
153 |    "metadata": {},
154 |    "outputs": [],
155 |    "source": [
156 |     "# List all available reformulation methods\n",
157 |     "from querygym.core.base import BaseReformulator\n",
158 |     "\n",
159 |     "methods = BaseReformulator.list_methods()\n",
160 |     "print(f\"Available methods ({len(methods)}):\")\n",
161 |     "for method in sorted(methods):\n",
162 |     "    print(f\"  - {method}\")"
163 |    ]
164 |   },
165 |   {
166 |    "cell_type": "markdown",
167 |    "metadata": {},
168 |    "source": [
169 |     "## 7. Explore Prompt Bank"
170 |    ]
171 |   },
172 |   {
173 |    "cell_type": "code",
174 |    "execution_count": null,
175 |    "metadata": {},
176 |    "outputs": [],
177 |    "source": [
178 |     "# List all prompts in the prompt bank\n",
179 |     "prompt_bank = qg.PromptBank()\n",
180 |     "prompts = prompt_bank.list()\n",
181 |     "\n",
182 |     "print(f\"Available prompts ({len(prompts)}):\")\n",
183 |     "for prompt_id in sorted(prompts):\n",
184 |     "    meta = prompt_bank.get_meta(prompt_id)\n",
185 |     "    print(f\"  - {prompt_id}: {meta.get('description', 'No description')}\")"
186 |    ]
187 |   },
188 |   {
189 |    "cell_type": "markdown",
190 |    "metadata": {},
191 |    "source": [
192 |     "## 8. Check PySerini Installation"
193 |    ]
194 |   },
195 |   {
196 |    "cell_type": "code",
197 |    "execution_count": null,
198 |    "metadata": {},
199 |    "outputs": [],
200 |    "source": [
201 |     "# Verify PySerini is available\n",
202 |     "try:\n",
203 |     "    import pyserini\n",
204 |     "    print(f\"✓ PySerini is installed\")\n",
205 |     "    print(f\"  Version: {pyserini.__version__}\")\n",
206 |     "except ImportError:\n",
207 |     "    print(\"✗ PySerini is not available\")"
208 |    ]
209 |   },
210 |   {
211 |    "cell_type": "markdown",
212 |    "metadata": {},
213 |    "source": [
214 |     "## 9. Check PyTorch & GPU"
215 |    ]
216 |   },
217 |   {
218 |    "cell_type": "code",
219 |    "execution_count": null,
220 |    "metadata": {},
221 |    "outputs": [],
222 |    "source": [
223 |     "# Check PyTorch and GPU availability\n",
224 |     "import torch\n",
225 |     "\n",
226 |     "print(f\"PyTorch version: {torch.__version__}\")\n",
227 |     "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
228 |     "\n",
229 |     "if torch.cuda.is_available():\n",
230 |     "    print(f\"CUDA version: {torch.version.cuda}\")\n",
231 |     "    print(f\"GPU count: {torch.cuda.device_count()}\")\n",
232 |     "    print(f\"GPU name: {torch.cuda.get_device_name(0)}\")\n",
233 |     "else:\n",
234 |     "    print(\"Running on CPU (no GPU detected)\")"
235 |    ]
236 |   },
237 |   {
238 |    "cell_type": "markdown",
239 |    "metadata": {},
240 |    "source": [
241 |     "## Next Steps\n",
242 |     "\n",
243 |     "- Load your own queries from `data/` directory\n",
244 |     "- Try different reformulation methods\n",
245 |     "- Experiment with different LLM models and parameters\n",
246 |     "- Check the [documentation](https://querygym.readthedocs.io/) for more examples"
247 |    ]
248 |   }
249 |  ],
250 |  "metadata": {
251 |   "kernelspec": {
252 |    "display_name": "Python 3",
253 |    "language": "python",
254 |    "name": "python3"
255 |   },
256 |   "language_info": {
257 |    "codemirror_mode": {
258 |     "name": "ipython",
259 |     "version": 3
260 |    },
261 |    "file_extension": ".py",
262 |    "mimetype": "text/x-python",
263 |    "name": "python",
264 |    "nbconvert_exporter": "python",
265 |    "pygments_lexer": "ipython3",
266 |    "version": "3.10.0"
267 |   }
268 |  },
269 |  "nbformat": 4,
270 |  "nbformat_minor": 4
271 | }
272 | 


--------------------------------------------------------------------------------
/querygym/adapters/pyserini_adapter.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | """
  3 | Pyserini Searcher Adapter for querygym
  4 | 
  5 | This module provides an adapter that wraps Pyserini's LuceneSearcher and LuceneImpactSearcher
  6 | to implement the querygym BaseSearcher interface.
  7 | """
  8 | 
  9 | import os
 10 | import json
 11 | from typing import List, Dict, Any, Optional
 12 | from ..core.searcher import BaseSearcher, SearchHit, SearcherRegistry
 13 | 
 14 | 
 15 | class PyseriniSearcher(BaseSearcher):
 16 |     """
 17 |     Pyserini adapter implementing the BaseSearcher interface.
 18 |     
 19 |     This adapter wraps Pyserini's LuceneSearcher and LuceneImpactSearcher
 20 |     to provide a standardized interface for querygym.
 21 |     """
 22 |     
 23 |     def __init__(self, index: str, searcher_type: str = "bm25", 
 24 |                  encoder: Optional[str] = None, min_idf: float = 0,
 25 |                  k1: Optional[float] = None, b: Optional[float] = None,
 26 |                  rm3: bool = False, rocchio: bool = False,
 27 |                  rocchio_use_negative: bool = False,
 28 |                  answer_key: str = "contents", **kwargs):
 29 |         """
 30 |         Initialize Pyserini searcher.
 31 |         
 32 |         Args:
 33 |             index: Path to Lucene index or prebuilt index name
 34 |             searcher_type: Type of searcher ("bm25" or "impact")
 35 |             encoder: Encoder for Impact search
 36 |             min_idf: Minimum IDF for Impact search
 37 |             k1: BM25 k1 parameter
 38 |             b: BM25 b parameter
 39 |             rm3: Enable RM3 expansion
 40 |             rocchio: Enable Rocchio expansion
 41 |             rocchio_use_negative: Use negative feedback in Rocchio
 42 |             answer_key: Field name(s) to extract content from (pipe-separated)
 43 |             **kwargs: Additional arguments
 44 |         """
 45 |         try:
 46 |             from pyserini.search.lucene import LuceneSearcher, LuceneImpactSearcher
 47 |         except ImportError:
 48 |             raise ImportError("Pyserini is required for PyseriniSearcher. Install with: pip install pyserini")
 49 |         
 50 |         self.index = index
 51 |         self.searcher_type = searcher_type
 52 |         self.answer_key = answer_key
 53 |         self._searcher_info = {
 54 |             "name": "PyseriniSearcher",
 55 |             "type": searcher_type,
 56 |             "index": index,
 57 |             "encoder": encoder,
 58 |             "min_idf": min_idf,
 59 |             "k1": k1,
 60 |             "b": b,
 61 |             "rm3": rm3,
 62 |             "rocchio": rocchio,
 63 |             "rocchio_use_negative": rocchio_use_negative,
 64 |             "answer_key": answer_key
 65 |         }
 66 |         
 67 |         # Initialize searcher based on type
 68 |         if searcher_type == "impact":
 69 |             if os.path.exists(index):
 70 |                 self.searcher = LuceneImpactSearcher(index, encoder, min_idf)
 71 |             else:
 72 |                 self.searcher = LuceneImpactSearcher.from_prebuilt_index(index, encoder, min_idf)
 73 |         else:  # Default to BM25
 74 |             if os.path.exists(index):
 75 |                 self.searcher = LuceneSearcher(index)
 76 |             else:
 77 |                 self.searcher = LuceneSearcher.from_prebuilt_index(index)
 78 |             
 79 |             # Set BM25 parameters if specified
 80 |             if k1 is not None and b is not None:
 81 |                 self.searcher.set_bm25(k1, b)
 82 |             else:
 83 |                 # Auto-set based on known indices
 84 |                 self._set_auto_bm25_params()
 85 |             
 86 |             # Add RM3 if requested
 87 |             if rm3:
 88 |                 self.searcher.set_rm3()
 89 |             
 90 |             # Add Rocchio if requested
 91 |             if rocchio:
 92 |                 if rocchio_use_negative:
 93 |                     self.searcher.set_rocchio(gamma=0.15, use_negative=True)
 94 |                 else:
 95 |                     self.searcher.set_rocchio()
 96 |     
 97 |     def _set_auto_bm25_params(self):
 98 |         """Set BM25 parameters based on index type."""
 99 |         if 'msmarco-passage' in self.index:
100 |             self.searcher.set_bm25(0.82, 0.68)
101 |         elif 'msmarco-doc' in self.index:
102 |             self.searcher.set_bm25(4.46, 0.82)
103 |     
104 |     def search(self, query: str, k: int = 10, **kwargs) -> List[SearchHit]:
105 |         """Search for documents using a single query."""
106 |         hits = self.searcher.search(query, k)
107 |         return self._process_hits(hits)
108 |     
109 |     def batch_search(self, queries: List[str], k: int = 10, 
110 |                     num_threads: int = 1, **kwargs) -> List[List[SearchHit]]:
111 |         """Search for documents using multiple queries in batch."""
112 |         pseudo_batch_topic_ids = [str(idx) for idx, _ in enumerate(queries)]
113 |         results = self.searcher.batch_search(
114 |             queries, pseudo_batch_topic_ids, k, num_threads
115 |         )
116 |         
117 |         # Convert results to list in original order
118 |         batch_results = [results[id_] for id_ in pseudo_batch_topic_ids]
119 |         
120 |         # Process each result
121 |         processed_results = []
122 |         for hits in batch_results:
123 |             processed_results.append(self._process_hits(hits))
124 |         
125 |         return processed_results
126 |     
127 |     def _process_hits(self, hits) -> List[SearchHit]:
128 |         """Process hits and extract content."""
129 |         search_hits = []
130 |         
131 |         for hit in hits:
132 |             try:
133 |                 # Extract content using hit.lucene_document.get('raw')
134 |                 raw_content = hit.lucene_document.get('raw')
135 |                 if raw_content:
136 |                     raw_df = json.loads(raw_content)
137 |                     text_list = [raw_df[k] for k in self.answer_key.split("|") if raw_df.get(k)]
138 |                     content = "\t".join(text_list)
139 |                 else:
140 |                     # Fallback: try to get content from other fields
141 |                     content = getattr(hit, 'contents', '') or getattr(hit, 'text', '') or str(hit.docid)
142 |                 
143 |                 search_hit = SearchHit(
144 |                     docid=hit.docid,
145 |                     score=hit.score,
146 |                     content=content,
147 |                     metadata={
148 |                         "raw_content": raw_content,
149 |                         "answer_key": self.answer_key
150 |                     }
151 |                 )
152 |                 search_hits.append(search_hit)
153 |                 
154 |             except (json.JSONDecodeError, KeyError, AttributeError) as e:
155 |                 # Fallback for malformed or missing content
156 |                 search_hit = SearchHit(
157 |                     docid=getattr(hit, 'docid', 'unknown'),
158 |                     score=getattr(hit, 'score', 0.0),
159 |                     content=str(getattr(hit, 'docid', 'unknown')),
160 |                     metadata={
161 |                         "error": str(e),
162 |                         "fallback": True
163 |                     }
164 |                 )
165 |                 search_hits.append(search_hit)
166 |         
167 |         return search_hits
168 |     
169 |     def get_searcher_info(self) -> Dict[str, Any]:
170 |         """Get information about the searcher configuration."""
171 |         return self._searcher_info.copy()
172 |     
173 |     def configure(self, **kwargs) -> None:
174 |         """Configure searcher parameters."""
175 |         # Update searcher info with new configuration
176 |         self._searcher_info.update(kwargs)
177 |         
178 |         # Apply configuration if possible
179 |         if "k1" in kwargs and "b" in kwargs:
180 |             self.searcher.set_bm25(kwargs["k1"], kwargs["b"])
181 |         
182 |         if "rm3" in kwargs and kwargs["rm3"]:
183 |             self.searcher.set_rm3()
184 |         
185 |         if "rocchio" in kwargs and kwargs["rocchio"]:
186 |             if kwargs.get("rocchio_use_negative", False):
187 |                 self.searcher.set_rocchio(gamma=0.15, use_negative=True)
188 |             else:
189 |                 self.searcher.set_rocchio()
190 | 
191 | 
192 | # Register the searcher
193 | SearcherRegistry.register("pyserini", PyseriniSearcher)
194 | 


--------------------------------------------------------------------------------
/querygym/methods/qa_expand.py:
--------------------------------------------------------------------------------
  1 | from __future__ import annotations
  2 | import json
  3 | import re
  4 | from ..core.base import BaseReformulator, QueryItem, ReformulationResult
  5 | from ..core.registry import register_method
  6 | 
  7 | def parse_llm_json(text: str) -> dict:
  8 |     """
  9 |     Robust JSON parser for LLM outputs that handles common formatting issues:
 10 |     - Markdown code blocks (```json, ```)
 11 |     - Escaped quotes (\")
 12 |     - Trailing commas
 13 |     - Single quotes instead of double quotes
 14 |     - Incomplete/truncated JSON
 15 |     - Extracts key-value pairs even from malformed JSON
 16 |     """
 17 |     if not text or not text.strip():
 18 |         return {}
 19 |     
 20 |     original_text = text
 21 |     text = text.strip()
 22 |     
 23 |     # Remove markdown code blocks
 24 |     if text.startswith("```"):
 25 |         first_newline = text.find('\n')
 26 |         if first_newline != -1:
 27 |             text = text[first_newline+1:]
 28 |         if text.endswith("```"):
 29 |             text = text[:-3]
 30 |         text = text.strip()
 31 |     
 32 |     # Unescape common escape sequences that appear in raw output
 33 |     text = text.replace('\\"', '"').replace('\\n', '\n').replace('\\t', '\t')
 34 |     
 35 |     # Try standard JSON parsing first
 36 |     try:
 37 |         return json.loads(text)
 38 |     except json.JSONDecodeError:
 39 |         pass
 40 |     
 41 |     # Remove trailing commas before closing braces/brackets
 42 |     cleaned = re.sub(r',(\s*[}\]])', r'\1', text)
 43 |     try:
 44 |         return json.loads(cleaned)
 45 |     except json.JSONDecodeError:
 46 |         pass
 47 |     
 48 |     # Replace single quotes with double quotes
 49 |     try:
 50 |         return json.loads(text.replace("'", '"'))
 51 |     except json.JSONDecodeError:
 52 |         pass
 53 |     
 54 |     # Try to fix incomplete JSON by adding missing closing quotes and braces
 55 |     if '{' in text:
 56 |         # Find the last complete value and try to close it
 57 |         attempt = text.rstrip()
 58 |         
 59 |         # If ends mid-string (no closing quote), add closing quote
 60 |         if attempt.count('"') % 2 == 1:
 61 |             attempt += '"'
 62 |         
 63 |         # Add missing closing brace if needed
 64 |         if not attempt.endswith('}'):
 65 |             attempt += '}'
 66 |         
 67 |         try:
 68 |             return json.loads(attempt)
 69 |         except json.JSONDecodeError:
 70 |             pass
 71 |     
 72 |     # Fallback: Extract key-value pairs using regex
 73 |     # Looks for patterns like "key": "value" or 'key': 'value'
 74 |     result = {}
 75 |     
 76 |     # Pattern to match JSON key-value pairs
 77 |     patterns = [
 78 |         r'["\'](\w+)["\']\s*:\s*["\']([^"\']*)["\']',  # Standard JSON
 79 |         r'(\w+)\s*:\s*["\']([^"\']*)["\']',             # Unquoted keys
 80 |         r'["\'](\w+)["\']\s*:\s*([^,}\]]+)',            # Values without quotes
 81 |     ]
 82 |     
 83 |     for pattern in patterns:
 84 |         matches = re.findall(pattern, text, re.DOTALL)
 85 |         for key, value in matches:
 86 |             if key and value:
 87 |                 result[key] = value.strip()
 88 |     
 89 |     # If we found any key-value pairs, return them
 90 |     if result:
 91 |         return result
 92 |     
 93 |     # Absolute fallback: return empty dict
 94 |     return {}
 95 | 
 96 | @register_method("qa_expand")
 97 | class QAExpand(BaseReformulator):
 98 |     """
 99 |     QA-Expand Pipeline (from paper):
100 |     1. Generate 3 sub-questions from original query (LLM call #1)
101 |     2. Generate 3 pseudo-answers from the 3 sub-questions (LLM call #2) - NO contexts
102 |     3. Filter and refine answers based on relevance to query (LLM call #3)
103 |     4. Expand query: (Q × 3) + refined_answers
104 |     """
105 |     VERSION = "1.0"
106 |     REQUIRES_CONTEXT = False  # No contexts needed - pure generative expansion
107 |     CONCATENATION_STRATEGY = "query_repeat_plus_generated"  # (Q × 3) + refined_answers
108 |     DEFAULT_QUERY_REPEATS = 3  # K=3 (fixed)
109 | 
110 |     def reformulate(self, q: QueryItem, contexts=None) -> ReformulationResult:
111 |         # K is fixed at 3 for QA-Expand
112 |         k = 3
113 |         
114 |         # Get configurable parameters
115 |         max_tokens = self.cfg.params.get("max_tokens", self.cfg.llm.get("max_tokens", 256))
116 |         default_temp = self.cfg.llm.get("temperature", 0.8)
117 |         
118 |         # Temperature settings (can override per-step, defaults to global temperature)
119 |         temp_subq = self.cfg.params.get("temperature_subq", default_temp)
120 |         temp_answer = self.cfg.params.get("temperature_answer", default_temp)
121 |         temp_refine = self.cfg.params.get("temperature_refine", default_temp)
122 |         
123 |         # Prompt IDs (customizable for different datasets)
124 |         prompt_subq = self.cfg.params.get("prompt_subq", "qa_expand.subq.v1")
125 |         prompt_answer = self.cfg.params.get("prompt_answer", "qa_expand.answer.v1")
126 |         prompt_refine = self.cfg.params.get("prompt_refine", "qa_expand.refine.v1")
127 |         
128 |         # Step 1: Generate 3 sub-questions (LLM call #1)
129 |         # Input: {query}
130 |         # Output: Raw text with 3 questions
131 |         msgs = self.prompts.render(prompt_subq, query=q.text)
132 |         subqs_raw = self.llm.chat(msgs, temperature=temp_subq, max_tokens=max_tokens)
133 |         
134 |         # Format questions as JSON for next step
135 |         # Try to parse if LLM already returned JSON, otherwise split by newlines
136 |         questions_dict = parse_llm_json(subqs_raw)
137 |         
138 |         if questions_dict and any(k.startswith('question') for k in questions_dict.keys()):
139 |             # LLM returned JSON format already
140 |             questions_json = json.dumps(questions_dict)
141 |         else:
142 |             # Parse from text - split by newlines and clean up
143 |             lines = [line.strip() for line in subqs_raw.split('\n') if line.strip()]
144 |             # Remove common prefixes like "1.", "Q1:", "-", etc.
145 |             cleaned_lines = []
146 |             for line in lines:
147 |                 line = re.sub(r'^[\d\.\-\*\)]+\s*', '', line)  # Remove "1.", "1)", "-", "*"
148 |                 line = re.sub(r'^[Qq]\d+[\:\.]?\s*', '', line)  # Remove "Q1:", "q1."
149 |                 if line:
150 |                     cleaned_lines.append(line)
151 |             
152 |             # Take first 3 questions
153 |             questions_json = json.dumps({
154 |                 "question1": cleaned_lines[0] if len(cleaned_lines) > 0 else "",
155 |                 "question2": cleaned_lines[1] if len(cleaned_lines) > 1 else "",
156 |                 "question3": cleaned_lines[2] if len(cleaned_lines) > 2 else ""
157 |             })
158 |         
159 |         # Step 2: Generate 3 pseudo-answers from LLM's internal knowledge (LLM call #2)
160 |         # Input: {questions} (JSON format)
161 |         # Output: JSON with {"answer1": "...", "answer2": "...", "answer3": "..."}
162 |         # NOTE: No contexts used - LLM generates answers from its own knowledge
163 |         msgs = self.prompts.render(prompt_answer, questions=questions_json)
164 |         answers_json = self.llm.chat(msgs, temperature=temp_answer, max_tokens=max_tokens)
165 |         
166 |         # Step 3: Filter and refine answers (LLM call #3)
167 |         # Input: {query}, {answers} (JSON format)
168 |         # Output: Refined JSON with relevant answers only
169 |         msgs = self.prompts.render(prompt_refine, query=q.text, answers=answers_json)
170 |         refined_answers_json = self.llm.chat(msgs, temperature=temp_refine, max_tokens=max_tokens)
171 |         
172 |         # Parse refined JSON and extract answer values for concatenation
173 |         refined_dict = parse_llm_json(refined_answers_json)
174 |         if refined_dict:
175 |             # Extract answer values in order (answer1, answer2, answer3)
176 |             answer_values = [str(v) for k, v in sorted(refined_dict.items()) if v and str(v).strip()]
177 |             refined_text = " ".join(answer_values)
178 |         else:
179 |             # Fallback if JSON parsing fails - use raw text
180 |             refined_text = refined_answers_json
181 |         
182 |         # Step 4: Expand query: (Q × 3) + refined_answers
183 |         reformulated = self.concatenate_result(q.text, refined_text, query_repeats=k)
184 |         
185 |         return ReformulationResult(
186 |             q.qid, 
187 |             q.text, 
188 |             reformulated, 
189 |             metadata={
190 |                 "subquestions_raw": subqs_raw,
191 |                 "questions_json": questions_json,
192 |                 "answers_json": answers_json,
193 |                 "refined_answers_json": refined_answers_json,
194 |                 "refined_text": refined_text,
195 |                 "prompts_used": {
196 |                     "subq": prompt_subq,
197 |                     "answer": prompt_answer,
198 |                     "refine": prompt_refine
199 |                 }
200 |             }
201 |         )
202 | 


--------------------------------------------------------------------------------
/_layouts/default.html:
--------------------------------------------------------------------------------
  1 | <!DOCTYPE html>
  2 | <html lang="en-US">
  3 |   <head>
  4 |     <meta charset="UTF-8">
  5 |     <meta http-equiv="X-UA-Compatible" content="IE=edge">
  6 |     <meta name="viewport" content="width=device-width, initial-scale=1">
  7 |     <title>{{ page.title | default: site.title }}</title>
  8 |     <meta name="description" content="{{ page.description | default: site.description }}">
  9 |     <link rel="preconnect" href="https://fonts.googleapis.com">
 10 |     <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
 11 |     <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap" rel="stylesheet">
 12 |     <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css">
 13 |     <style>
 14 |       :root {
 15 |         --primary-color: #667eea;
 16 |         --primary-dark: #764ba2;
 17 |         --bg-body: #ffffff;
 18 |         --bg-code: #f6f8fa;
 19 |         --text-primary: #24292f;
 20 |         --text-secondary: #57606a;
 21 |         --border-color: #d0d7de;
 22 |         --header-gradient: linear-gradient(to right, #667eea 0%, #764ba2 100%);
 23 |         --link-color: #0969da;
 24 |       }
 25 | 
 26 |       [data-theme="dark"] {
 27 |         --primary-color: #818cf8;
 28 |         --primary-dark: #a78bfa;
 29 |         --bg-body: #0d1117;
 30 |         --bg-code: #161b22;
 31 |         --text-primary: #c9d1d9;
 32 |         --text-secondary: #8b949e;
 33 |         --border-color: #30363d;
 34 |         --header-gradient: linear-gradient(to right, #4338ca 0%, #6b21a8 100%);
 35 |         --link-color: #58a6ff;
 36 |       }
 37 | 
 38 |       * {
 39 |         box-sizing: border-box;
 40 |       }
 41 | 
 42 |       body {
 43 |         font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Helvetica, Arial, sans-serif;
 44 |         font-size: 16px;
 45 |         line-height: 1.6;
 46 |         color: var(--text-primary);
 47 |         background-color: var(--bg-body);
 48 |         margin: 0;
 49 |         padding: 0;
 50 |         transition: background-color 0.3s ease, color 0.3s ease;
 51 |       }
 52 | 
 53 |       a {
 54 |         color: var(--link-color);
 55 |         text-decoration: none;
 56 |       }
 57 | 
 58 |       a:hover {
 59 |         text-decoration: underline;
 60 |       }
 61 | 
 62 |       .page-header {
 63 |         background: var(--header-gradient);
 64 |         padding: 2rem 1rem;
 65 |         text-align: center;
 66 |         color: #fff;
 67 |         position: relative;
 68 |       }
 69 | 
 70 |       .header-controls {
 71 |         position: absolute;
 72 |         top: 1rem;
 73 |         right: 1rem;
 74 |         display: flex;
 75 |         gap: 0.5rem;
 76 |       }
 77 | 
 78 |       .header-btn {
 79 |         display: inline-flex;
 80 |         align-items: center;
 81 |         justify-content: center;
 82 |         width: 40px;
 83 |         height: 40px;
 84 |         background: rgba(255, 255, 255, 0.15);
 85 |         border: 1px solid rgba(255, 255, 255, 0.3);
 86 |         border-radius: 6px;
 87 |         color: white;
 88 |         cursor: pointer;
 89 |         transition: all 0.2s ease;
 90 |       }
 91 | 
 92 |       .header-btn:hover {
 93 |         background: rgba(255, 255, 255, 0.25);
 94 |         text-decoration: none;
 95 |       }
 96 | 
 97 |       .project-name {
 98 |         margin: 0;
 99 |         font-size: 2rem;
100 |         font-weight: 700;
101 |       }
102 | 
103 |       .project-tagline {
104 |         margin: 0.5rem 0 0;
105 |         opacity: 0.9;
106 |         font-size: 1.1rem;
107 |         font-weight: 400;
108 |       }
109 | 
110 |       .main-content {
111 |         max-width: 900px;
112 |         margin: 0 auto;
113 |         padding: 2rem 1rem;
114 |       }
115 | 
116 |       .main-content h1,
117 |       .main-content h2,
118 |       .main-content h3,
119 |       .main-content h4,
120 |       .main-content h5,
121 |       .main-content h6 {
122 |         color: var(--text-primary);
123 |         margin-top: 2rem;
124 |         margin-bottom: 1rem;
125 |         font-weight: 600;
126 |       }
127 | 
128 |       .main-content h1 { font-size: 2rem; border-bottom: 1px solid var(--border-color); padding-bottom: 0.5rem; }
129 |       .main-content h2 { font-size: 1.5rem; border-bottom: 1px solid var(--border-color); padding-bottom: 0.3rem; }
130 |       .main-content h3 { font-size: 1.25rem; }
131 | 
132 |       .main-content p {
133 |         margin-bottom: 1rem;
134 |       }
135 | 
136 |       .main-content ul,
137 |       .main-content ol {
138 |         margin-bottom: 1rem;
139 |         padding-left: 2rem;
140 |       }
141 | 
142 |       .main-content li {
143 |         margin-bottom: 0.5rem;
144 |       }
145 | 
146 |       .main-content code {
147 |         background: var(--bg-code);
148 |         padding: 0.2rem 0.4rem;
149 |         border-radius: 4px;
150 |         font-size: 0.9em;
151 |         font-family: 'SFMono-Regular', Consolas, 'Liberation Mono', Menlo, monospace;
152 |       }
153 | 
154 |       .main-content pre {
155 |         background: var(--bg-code);
156 |         padding: 1rem;
157 |         border-radius: 6px;
158 |         overflow-x: auto;
159 |         margin-bottom: 1rem;
160 |         border: 1px solid var(--border-color);
161 |       }
162 | 
163 |       .main-content pre code {
164 |         background: transparent;
165 |         padding: 0;
166 |       }
167 | 
168 |       .main-content img {
169 |         max-width: 100%;
170 |         height: auto;
171 |       }
172 | 
173 |       .main-content blockquote {
174 |         margin: 1rem 0;
175 |         padding: 0.5rem 1rem;
176 |         border-left: 4px solid var(--primary-color);
177 |         background: var(--bg-code);
178 |         color: var(--text-secondary);
179 |       }
180 | 
181 |       .main-content hr {
182 |         border: none;
183 |         border-top: 1px solid var(--border-color);
184 |         margin: 2rem 0;
185 |       }
186 | 
187 |       .main-content table {
188 |         width: 100%;
189 |         border-collapse: collapse;
190 |         margin-bottom: 1rem;
191 |       }
192 | 
193 |       .main-content th,
194 |       .main-content td {
195 |         padding: 0.75rem;
196 |         border: 1px solid var(--border-color);
197 |         text-align: left;
198 |       }
199 | 
200 |       .main-content th {
201 |         background: var(--bg-code);
202 |         font-weight: 600;
203 |       }
204 | 
205 |       .site-footer {
206 |         border-top: 1px solid var(--border-color);
207 |         margin-top: 2rem;
208 |         padding: 2rem 1rem;
209 |         text-align: center;
210 |         color: var(--text-secondary);
211 |         font-size: 0.875rem;
212 |       }
213 | 
214 |       @media screen and (min-width: 64em) {
215 |         .page-header {
216 |           padding: 3rem 2rem;
217 |         }
218 |         .project-name {
219 |           font-size: 2.5rem;
220 |         }
221 |         .main-content {
222 |           padding: 2rem;
223 |         }
224 |       }
225 |     </style>
226 |   </head>
227 |   <body>
228 |     <header class="page-header">
229 |       <div class="header-controls">
230 |         <a href="https://github.com/ls3-lab/QueryGym" class="header-btn" title="View on GitHub">
231 |           <i class="fab fa-github"></i>
232 |         </a>
233 |         <button class="header-btn" id="themeToggle" title="Toggle dark mode">
234 |           <i class="fas fa-moon"></i>
235 |         </button>
236 |       </div>
237 |       <h1 class="project-name">{{ site.title | default: "QueryGym" }}</h1>
238 |       <h2 class="project-tagline">{{ site.description | default: "A lightweight, reproducible toolkit for LLM-based query reformulation" }}</h2>
239 |     </header>
240 | 
241 |     <main class="main-content">
242 |       {{ content }}
243 |     </main>
244 | 
245 |     <footer class="site-footer">
246 |       <span>Made with ❤️ by the QueryGym Team</span>
247 |       <br>
248 |       <span style="font-size: 0.8rem; margin-top: 0.5rem; display: inline-block;">© 2024-2025 QueryGym Contributors · <a href="https://github.com/ls3-lab/QueryGym">GitHub</a> · <a href="https://opensource.org/licenses/Apache-2.0">Apache 2.0</a></span>
249 |     </footer>
250 | 
251 |     <script>
252 |       // Dark mode toggle
253 |       const themeToggle = document.getElementById('themeToggle');
254 |       const themeIcon = themeToggle.querySelector('i');
255 |       
256 |       // Check for saved theme preference or default to light
257 |       const savedTheme = localStorage.getItem('querygym-theme') || 'light';
258 |       if (savedTheme === 'dark') {
259 |         document.documentElement.setAttribute('data-theme', 'dark');
260 |         themeIcon.classList.remove('fa-moon');
261 |         themeIcon.classList.add('fa-sun');
262 |       }
263 |       
264 |       themeToggle.addEventListener('click', () => {
265 |         const currentTheme = document.documentElement.getAttribute('data-theme');
266 |         const newTheme = currentTheme === 'dark' ? 'light' : 'dark';
267 |         
268 |         document.documentElement.setAttribute('data-theme', newTheme);
269 |         localStorage.setItem('querygym-theme', newTheme);
270 |         
271 |         if (newTheme === 'dark') {
272 |           themeIcon.classList.remove('fa-moon');
273 |           themeIcon.classList.add('fa-sun');
274 |         } else {
275 |           themeIcon.classList.remove('fa-sun');
276 |           themeIcon.classList.add('fa-moon');
277 |         }
278 |       });
279 |     </script>
280 |   </body>
281 | </html>
282 | 


--------------------------------------------------------------------------------
/querygym/methods/csqe.py:
--------------------------------------------------------------------------------
  1 | from __future__ import annotations
  2 | from ..core.base import BaseReformulator, QueryItem, ReformulationResult
  3 | from ..core.registry import register_method
  4 | from typing import List, Optional, Dict, Any
  5 | import re
  6 | 
  7 | @register_method("csqe")
  8 | class CSQE(BaseReformulator):
  9 |     """
 10 |     CSQE (Context-based Sentence-level Query Expansion) method.
 11 |     
 12 |     Combines KEQE (knowledge-based) and CSQE (context-based) expansions:
 13 |     1. Generates N KEQE passages from LLM knowledge (no contexts) - default N=2
 14 |     2. Generates N CSQE responses with retrieved contexts (few-shot) - default N=2
 15 |     3. Extracts key sentences from CSQE responses
 16 |     4. Concatenates: query × N + keqe_passages + csqe_sentences (newline-separated, lowercased)
 17 |     
 18 |     Default: N=2 for both, totaling 4 generations (2 KEQE + 2 CSQE).
 19 |     Query is repeated N times (equal to number of KEQE expansions).
 20 |     """
 21 |     VERSION = "1.0"
 22 |     REQUIRES_CONTEXT = True  # Needs retrieved contexts for CSQE expansion
 23 | 
 24 |     def reformulate(self, q: QueryItem, contexts=None) -> ReformulationResult:
 25 |         """Reformulate query using CSQE method (KEQE + CSQE)."""
 26 |         # Use provided contexts (from batch retrieval)
 27 |         ctxs = contexts or []
 28 | 
 29 |         # Prepare contexts for CSQE prompting
 30 |         retrieval_k = int(self.cfg.params.get("retrieval_k", 10))
 31 |         prompt_ctxs = ctxs[:retrieval_k]
 32 |         contexts_blob = "\n".join([f"{i+1}. {psg}" for i, psg in enumerate(prompt_ctxs)])
 33 |         
 34 |         # Get generation parameters - N=2 for both by default (paper specification)
 35 |         n_expansions = int(self.cfg.params.get("gen_num", 2))  # Number of expansions for both KEQE and CSQE (default: 2)
 36 |         max_tokens = int(self.cfg.llm.get("max_tokens", 1024))
 37 |         temperature = float(self.cfg.llm.get("temperature", 1.0))
 38 | 
 39 |         # Step 1: KEQE expansion - Generate N passages from knowledge (no contexts) in one call
 40 |         msgs_keqe = self.prompts.render("keqe.v1", query=q.text)
 41 |         # Use n parameter to generate all passages in one call (matching new_baselines.py)
 42 |         resp_keqe = self.llm.client.chat.completions.create(
 43 |             model=self.llm.model,
 44 |             messages=msgs_keqe,
 45 |             n=n_expansions,
 46 |             temperature=temperature,
 47 |             max_tokens=max_tokens
 48 |         )
 49 |         keqe_passages = [
 50 |             choice.message.content.strip().strip('"').strip("'") or ""
 51 |             for choice in resp_keqe.choices
 52 |         ]
 53 | 
 54 |         # Step 2: CSQE expansion - Generate N responses with contexts (few-shot) in one call
 55 |         msgs_csqe = self.prompts.render("csqe.v1", query=q.text, contexts=contexts_blob)
 56 |         # Use n parameter to generate all responses in one call (matching new_baselines.py)
 57 |         resp_csqe = self.llm.client.chat.completions.create(
 58 |             model=self.llm.model,
 59 |             messages=msgs_csqe,
 60 |             n=n_expansions,
 61 |             temperature=temperature,
 62 |             max_tokens=max_tokens
 63 |         )
 64 |         csqe_responses = [
 65 |             choice.message.content or ""
 66 |             for choice in resp_csqe.choices
 67 |         ]
 68 |         # Step 3: Extract key sentences from CSQE responses
 69 |         # Pattern matches sentences in double quotes (from new_baselines.py)
 70 |         csqe_sentences: List[str] = []
 71 |         for response in csqe_responses:
 72 |             extracted = self._extract_key_sentences(response)
 73 |             csqe_sentences.append(extracted)
 74 | 
 75 |         # Step 4: Concatenate directly (query × N + keqe_passages + csqe_sentences)
 76 |         # Format: query repeated N times (N = number of KEQE passages),
 77 |         # then all KEQE passages, then all extracted CSQE sentences, space-separated and lowercased
 78 |         parts = [q.text] * n_expansions + keqe_passages + csqe_sentences
 79 |         reformulated = " ".join(parts).lower().strip()
 80 |         # Clean up quotes
 81 |         reformulated = reformulated.strip('"').strip("'")
 82 | 
 83 |         return ReformulationResult(
 84 |             q.qid,
 85 |             q.text,
 86 |             reformulated,
 87 |             metadata={
 88 |                 "keqe_passages": keqe_passages,
 89 |                 "csqe_responses": csqe_responses,
 90 |                 "csqe_sentences": csqe_sentences,
 91 |                 "gen_num": n_expansions,
 92 |                 "total_generations": n_expansions * 2,  # KEQE + CSQE
 93 |                 "used_ctx": len(prompt_ctxs),
 94 |             },
 95 |         )
 96 | 
 97 |     def _extract_key_sentences(self, response: str) -> str:
 98 |         """
 99 |         Extract key sentences from CSQE response.
100 |         
101 |         Primary: Uses regex to find sentences in double quotes (expected format).
102 |         Fallback: If no quotes found, extracts content after numbered document markers (1., 2., etc.)
103 |         """
104 |         # Primary strategy: Extract sentences in double quotes
105 |         pattern = r'"([^"]*)"'
106 |         sentences = re.findall(pattern, response)
107 |         
108 |         if sentences:
109 |             # Found quoted sentences - join and return
110 |             joint_sentence = " ".join(sentences)
111 |             return joint_sentence
112 |         
113 |         # Fallback: LLM didn't use quotes, extract from numbered document format
114 |         # Pattern: Look for content after "1.", "2.", etc. (with optional trailing colon)
115 |         # Remove "Relevant Documents:" header first
116 |         cleaned = re.sub(r'^Relevant Documents?:?\s*\n?', '', response, flags=re.IGNORECASE | re.MULTILINE)
117 |         
118 |         # Extract content after numbered markers (1., 2., etc.)
119 |         doc_pattern = r'\d+[\.:]\s*(.+?)(?=\d+[\.:]|$)'
120 |         doc_content = re.findall(doc_pattern, cleaned, re.DOTALL)
121 |         
122 |         if doc_content:
123 |             # Clean up each extracted piece
124 |             extracted = []
125 |             for content in doc_content:
126 |                 # Remove extra whitespace and newlines, limit length
127 |                 cleaned_content = " ".join(content.split())
128 |                 if cleaned_content:
129 |                     extracted.append(cleaned_content)
130 |             if extracted:
131 |                 return " ".join(extracted)
132 |         
133 |         # If still nothing found, return empty string
134 |         return ""
135 | 
136 |     def _get_retrieval_params(self) -> Optional[Dict[str, Any]]:
137 |         """Get CSQE-specific retrieval parameters for batch retrieval.
138 |         
139 |         Returns retrieval parameters compatible with BaseSearcher interface.
140 |         Supports both searcher_type/searcher_kwargs format and legacy format.
141 |         """
142 |         # Check if searcher is provided directly (wrappers, adapter instances, etc.)
143 |         if "searcher" in self.cfg.params:
144 |             # Direct searcher instance provided
145 |             return {
146 |                 "searcher": self.cfg.params["searcher"],
147 |                 "k": int(self.cfg.params.get("retrieval_k", 10)),
148 |                 "threads": int(self.cfg.params.get("threads", 16)),
149 |             }
150 |         
151 |         # Check if using new searcher interface format with searcher_type
152 |         if "searcher_type" in self.cfg.params:
153 |             # New format: use searcher_type and searcher_kwargs
154 |             searcher_type = self.cfg.params.get("searcher_type", "pyserini")
155 |             searcher_kwargs = self.cfg.params.get("searcher_kwargs", {})
156 |             
157 |             # If index is provided in params but not in searcher_kwargs, add it
158 |             if "index" in self.cfg.params and "index" not in searcher_kwargs:
159 |                 searcher_kwargs["index"] = self.cfg.params["index"]
160 |             
161 |             return {
162 |                 "searcher_type": searcher_type,
163 |                 "searcher_kwargs": searcher_kwargs,
164 |                 "k": int(self.cfg.params.get("retrieval_k", 10)),
165 |                 "threads": int(self.cfg.params.get("threads", 16)),
166 |             }
167 |         
168 |         # Legacy format: convert old-style params to new format
169 |         index = self.cfg.params.get("index")
170 |         if not index:
171 |             return None
172 |         
173 |         # Build searcher_kwargs from legacy params
174 |         searcher_kwargs = {
175 |             "index": index,
176 |             "searcher_type": "impact" if self.cfg.params.get("impact", False) else "bm25",
177 |             "answer_key": self.cfg.params.get("answer_key", "contents"),
178 |         }
179 |         
180 |         # Add BM25 parameters if specified
181 |         k1 = self.cfg.params.get("k1")
182 |         b = self.cfg.params.get("b")
183 |         if k1 is not None and b is not None:
184 |             searcher_kwargs["k1"] = k1
185 |             searcher_kwargs["b"] = b
186 |         
187 |         # Add RM3 if requested
188 |         if self.cfg.params.get("rm3", False):
189 |             searcher_kwargs["rm3"] = True
190 |         
191 |         # Add Rocchio if requested
192 |         if self.cfg.params.get("rocchio", False):
193 |             searcher_kwargs["rocchio"] = True
194 |             if self.cfg.params.get("rocchio_use_negative", False):
195 |                 searcher_kwargs["rocchio_use_negative"] = True
196 |         
197 |         return {
198 |             "searcher_type": "pyserini",  # Default to pyserini for legacy format
199 |             "searcher_kwargs": searcher_kwargs,
200 |             "k": int(self.cfg.params.get("retrieval_k", 10)),
201 |             "threads": int(self.cfg.params.get("threads", 16)),
202 |         }
203 | 


--------------------------------------------------------------------------------