diff --git a/FACETING_STATUS.md b/FACETING_STATUS.md new file mode 100644 index 00000000..1e644c86 --- /dev/null +++ b/FACETING_STATUS.md @@ -0,0 +1,101 @@ +# ColBERT Faceting Feature Implementation Status + +## Overview + +This document summarizes the current status of the faceting and metadata filtering features in ColBERT. + +## Implemented Features + +1. **Metadata Support in Collection** + - Added metadata storage to Collection class + - Implemented JSONL loading/saving with metadata + - Created `get_metadata(pid)` method for retrieving passage metadata + +2. **Faceted Search** + - Added `facet_fields` parameter to search methods + - Implemented facet value computation on search results + - Added facet data to Ranking objects + +3. **Metadata Filtering** + - Added `facet_filters` parameter to search methods + - Implemented filtering based on exact matches + - Added support for list-based filters + - Added range filters with operators (`>=`, `<=`) + - Implemented combined filters across multiple fields + +4. **Server Integration** + - Modified server API to support facet_fields and facet_filters parameters + +## Test Infrastructure + +1. **Test Data Generation** + - Created `create_faceting_benchmark_data.py` for generating synthetic datasets + - Supports configurable metadata fields with different cardinalities + - Generates metadata with structured patterns for benchmarking + +2. **Test Framework** + - Added pytest option `--benchmark-data-dir` for specifying test data + - Created fixtures for loading pre-generated test data + - Configured directory structure for storing test data + +3. **Test Suites** + - `test_faceting_basic.py`: Basic functionality tests + - `test_faceting_data.py`: Tests for data loading + - `test_faceting_benchmark_mock.py`: Mock searcher tests + - `test_faceting_benchmark.py`: Framework for full benchmark tests + - `test_collection_metadata.py`: Tests for collection metadata functionality + - `test_faceting.py`: Tests for facet filtering with mock data + +## Documentation + +1. **README Updates** + - Added "Metadata and Faceted Search" section + - Included metadata storage format examples + - Added code examples for metadata filtering + - Added code examples for faceted search + - Included instructions for running benchmarks + +2. **Developer Documentation** + - Added test instructions and examples + - Documented benchmark data generation + +## Git Status + +1. **Current Branch**: `feature/add-faceting` + +2. **Modified Files**: + - `colbert/data/collection.py`: Added metadata support + - `colbert/data/ranking.py`: Added facet data storage + - `colbert/searcher.py`: Added facet filtering and computation + - `server.py`: Added facet support to API + - `README.md`: Added documentation + +3. **New Files**: + - `conftest.py`: Pytest configuration + - `pyproject.toml`: Project configuration + - `scripts/create_faceting_benchmark_data.py`: Benchmark data generator + - `tests/test_collection.py`: Collection tests + - `tests/test_collection_metadata.py`: Metadata tests + - `tests/test_faceting.py`: Faceting tests + - `tests/test_faceting_basic.py`: Basic faceting tests + - `tests/test_faceting_benchmark.py`: Benchmark framework + - `tests/test_faceting_benchmark_mock.py`: Mock benchmark tests + - `tests/test_faceting_data.py`: Data loading tests + - `tests/data/.gitignore`: Ignores synthetic test data + - `tests/data/faceting_benchmark/.gitkeep`: Preserves directory structure + +## Next Steps + +1. **Performance Optimization** + - Profile faceting performance with large datasets + - Optimize facet computation for high-cardinality fields + - Add caching for frequently used facet values + +2. **Additional Features** + - Add hierarchical faceting support + - Implement facet value counts for all results (not just top-k) + - Add pagination support for faceted results + +3. **Integration Testing** + - Test with real-world datasets + - Benchmark against other faceted search implementations \ No newline at end of file diff --git a/README.md b/README.md index ad000f06..9568833e 100644 --- a/README.md +++ b/README.md @@ -224,6 +224,73 @@ A sample query: http://localhost:8893/api/search?query=Who won the 2022 FIFA world cup&k=25 ``` +## Metadata and Faceted Search + +ColBERT supports metadata-based filtering and faceted search to enhance retrieval capabilities. This allows you to: + +1. Store and retrieve metadata alongside passages +2. Filter search results based on metadata fields +3. Compute facet counts for search results + +### Metadata Storage + +Metadata can be stored in JSONL format: + +```json +{"pid": 0, "passage": "Document text", "metadata": {"category": "science", "year": 2022}} +``` + +### Metadata Filtering + +Filter results based on metadata fields: + +```python +# Filter by exact match +results = searcher.search("quantum physics", facet_filters={"category": "science"}) + +# Filter by list of values +results = searcher.search("election analysis", facet_filters={"category": ["politics", "news"]}) + +# Filter by numeric range +results = searcher.search("recent discoveries", facet_filters={"year": ">=2020"}) + +# Combine multiple filters +results = searcher.search("recent research", facet_filters={ + "category": "science", + "year": ">=2020", + "author": ["Smith", "Jones"] +}) +``` + +### Faceted Search + +Compute facet counts for search results: + +```python +# Return facet counts along with search results +pids, ranks, scores, facets = searcher.search( + "climate change", + k=100, + facet_fields=["category", "year", "source"] +) + +# Access facet counts +print(f"Categories: {facets['category']}") # {'science': 42, 'news': 12, ...} +print(f"Years: {facets['year']}") # {'2022': 15, '2021': 25, ...} +``` + +### Benchmark + +To run the faceting benchmark: + +```bash +# Generate synthetic benchmark data +python scripts/create_faceting_benchmark_data.py --output-dir=tests/data/faceting_benchmark --num-passages=10000 + +# Run benchmark tests +python -m pytest tests/test_faceting_benchmark.py -v --benchmark-data-dir=tests/data/faceting_benchmark +``` + ## Branches ### Supported branches diff --git a/colbert/data/collection.py b/colbert/data/collection.py index d5efc943..96488f47 100644 --- a/colbert/data/collection.py +++ b/colbert/data/collection.py @@ -12,9 +12,17 @@ class Collection: - def __init__(self, path=None, data=None): + def __init__(self, path=None, data=None, metadata=None): self.path = path - self.data = data or self._load_file(path) + self.metadata = {} if metadata is None else metadata + + if data is not None: + self.data = data + elif path is not None: + self.data = self._load_file(path) + else: + # If both path and data are None, initialize with empty list + self.data = [] def __iter__(self): # TODO: If __data isn't there, stream from disk! @@ -23,6 +31,18 @@ def __iter__(self): def __getitem__(self, item): # TODO: Load from disk the first time this is called. Unless self.data is already not None. return self.data[item] + + def get_metadata(self, pid): + """ + Get metadata for a specific passage by PID. + + Args: + pid: The passage ID + + Returns: + Dict of metadata fields or empty dict if no metadata exists + """ + return self.metadata.get(pid, {}) def __len__(self): # TODO: Load here too. Basically, let's make data a property function and, on first call, either load or get __data. @@ -36,25 +56,83 @@ def _load_tsv(self, path): return load_collection(path) def _load_jsonl(self, path): - raise NotImplementedError() + """ + Load collection from JSONL file with metadata support. + Expected format for each line: + {"pid": 0, "passage": "text", "metadata": {"field1": "value1", "field2": "value2"}} + """ + import json + passages = [] + metadata = {} + + with open(path, 'r') as f: + for line_idx, line in enumerate(f): + if line_idx % (1000*1000) == 0: + print(f'{line_idx // 1000 // 1000}M', end=' ', flush=True) + + try: + record = json.loads(line.strip()) + + # Ensure required fields are present + pid = record.get('pid', line_idx) + passage = record.get('passage', '') + + # Add to collection + passages.append(passage) + + # Process metadata if present + if 'metadata' in record and isinstance(record['metadata'], dict): + metadata[pid] = record['metadata'] + + except json.JSONDecodeError: + # Skip malformed lines + print(f"Warning: Skipping malformed JSONL line {line_idx}") + except Exception as e: + print(f"Error processing line {line_idx}: {e}") + + # Store metadata in the object + self.metadata = metadata + + return passages def provenance(self): return self.path def toDict(self): - return {'provenance': self.provenance()} + return { + 'provenance': self.provenance(), + 'has_metadata': len(self.metadata) > 0 + } def save(self, new_path): - assert new_path.endswith('.tsv'), "TODO: Support .json[l] too." + assert new_path.endswith('.tsv') or new_path.endswith('.jsonl'), "Only .tsv and .jsonl formats are supported" assert not os.path.exists(new_path), new_path - with Run().open(new_path, 'w') as f: - # TODO: expects content to always be a string here; no separate title! - for pid, content in enumerate(self.data): - content = f'{pid}\t{content}\n' - f.write(content) - - return f.name + if new_path.endswith('.tsv'): + # Save in TSV format (without metadata) + with Run().open(new_path, 'w') as f: + for pid, content in enumerate(self.data): + content = f'{pid}\t{content}\n' + f.write(content) + + return f.name + else: + # Save in JSONL format with metadata + import json + with Run().open(new_path, 'w') as f: + for pid, content in enumerate(self.data): + record = { + 'pid': pid, + 'passage': content + } + + # Include metadata if it exists for this passage + if pid in self.metadata: + record['metadata'] = self.metadata[pid] + + f.write(json.dumps(record) + '\n') + + return f.name def enumerate(self, rank): for _, offset, passages in self.enumerate_batches(rank=rank): diff --git a/colbert/data/ranking.py b/colbert/data/ranking.py index c334a52f..ee394a93 100644 --- a/colbert/data/ranking.py +++ b/colbert/data/ranking.py @@ -23,15 +23,19 @@ def load_ranking(path): # works with annotated and un-annotated ranked lists class Ranking: - def __init__(self, path=None, data=None, metrics=None, provenance=None): + def __init__(self, path=None, data=None, metrics=None, provenance=None, metadata=None): self.__provenance = provenance or path or Provenance() self.data = self._prepare_data(data or self._load_file(path)) + self.metadata = metadata or {} def provenance(self): return self.__provenance def toDict(self): - return {'provenance': self.provenance()} + return { + 'provenance': self.provenance(), + 'metadata': self.metadata + } def _prepare_data(self, data): # TODO: Handle list of lists??? @@ -75,6 +79,11 @@ def save(self, new_path): d = {} d['metadata'] = get_metadata_only() d['provenance'] = self.provenance() + + # Include any custom metadata + if self.metadata: + d['custom_metadata'] = self.metadata + line = ujson.dumps(d, indent=4) f.write(line) diff --git a/colbert/searcher.py b/colbert/searcher.py index 8bc07c50..2785f554 100644 --- a/colbert/searcher.py +++ b/colbert/searcher.py @@ -62,36 +62,98 @@ def encode(self, text: TextQueries, full_length_search=False): return Q - def search(self, text: str, k=10, filter_fn=None, full_length_search=False, pids=None): + def search(self, text: str, k=10, filter_fn=None, full_length_search=False, pids=None, + facet_fields=None, facet_filters=None): + """ + Search for passages matching the given query. + + Args: + text: The search query text + k: Number of results to return + filter_fn: Function to filter results (used by base implementation) + full_length_search: Whether to use full-length search + pids: List of passage IDs to restrict search to + facet_fields: List of metadata fields to compute facet counts for + facet_filters: Dict of metadata fields to filter values, e.g. {"year": [2019, 2020], "category": "science"} + + Returns: + A tuple of (pids, ranks, scores) if facet_fields is None, otherwise + a tuple of (pids, ranks, scores, facets) where facets is a dict of field -> value -> count + """ Q = self.encode(text, full_length_search=full_length_search) - return self.dense_search(Q, k, filter_fn=filter_fn, pids=pids) + + # Create a filter function that combines the original filter with facet filtering + combined_filter_fn = self._create_facet_filter(filter_fn, facet_filters) + + # Get search results + pids, ranks, scores = self.dense_search(Q, k, filter_fn=combined_filter_fn, pids=pids) + + # If facet fields are requested, compute facet counts + if facet_fields: + facets = self._compute_facets(pids, facet_fields) + return pids, ranks, scores, facets + + return pids, ranks, scores - def search_all(self, queries: TextQueries, k=10, filter_fn=None, full_length_search=False, qid_to_pids=None): + def search_all(self, queries: TextQueries, k=10, filter_fn=None, full_length_search=False, + qid_to_pids=None, facet_fields=None, facet_filters=None): + """ + Search for multiple queries with optional faceting. + + Args: + queries: The search queries + k: Number of results to return per query + filter_fn: Function to filter results + full_length_search: Whether to use full-length search + qid_to_pids: Dict of query ID to list of passage IDs to restrict search to + facet_fields: List of metadata fields to compute facet counts for + facet_filters: Dict of metadata fields to filter values + + Returns: + A Ranking object, optionally with facet information + """ queries = Queries.cast(queries) queries_ = list(queries.values()) Q = self.encode(queries_, full_length_search=full_length_search) - return self._search_all_Q(queries, Q, k, filter_fn=filter_fn, qid_to_pids=qid_to_pids) + return self._search_all_Q(queries, Q, k, filter_fn=filter_fn, qid_to_pids=qid_to_pids, + facet_fields=facet_fields, facet_filters=facet_filters) - def _search_all_Q(self, queries, Q, k, filter_fn=None, qid_to_pids=None): + def _search_all_Q(self, queries, Q, k, filter_fn=None, qid_to_pids=None, + facet_fields=None, facet_filters=None): qids = list(queries.keys()) if qid_to_pids is None: qid_to_pids = {qid: None for qid in qids} - all_scored_pids = [ - list( - zip( - *self.dense_search( - Q[query_idx:query_idx+1], - k, filter_fn=filter_fn, - pids=qid_to_pids[qid] - ) + # Create the facet filter function once + combined_filter_fn = self._create_facet_filter(filter_fn, facet_filters) + + # Track facets separately if requested + all_facets = {} + all_scored_pids = [] + + for query_idx, qid in tqdm(enumerate(qids)): + # Use the combined filter function for search + if facet_fields: + # With facet computation + pids, ranks, scores, facets = self.search( + queries[qid], k, + filter_fn=combined_filter_fn, + pids=qid_to_pids[qid], + facet_fields=facet_fields ) - ) - for query_idx, qid in tqdm(enumerate(qids)) - ] + all_scored_pids.append(list(zip(pids, ranks, scores))) + all_facets[qid] = facets + else: + # Without facet computation + search_results = self.dense_search( + Q[query_idx:query_idx+1], + k, filter_fn=combined_filter_fn, + pids=qid_to_pids[qid] + ) + all_scored_pids.append(list(zip(*search_results))) data = {qid: val for qid, val in zip(queries.keys(), all_scored_pids)} @@ -100,8 +162,99 @@ def _search_all_Q(self, queries, Q, k, filter_fn=None, qid_to_pids=None): provenance.queries = queries.provenance() provenance.config = self.config.export() provenance.k = k + + if facet_fields: + # Add facet information to provenance + provenance.facets = True + # Create a ranking with facet data + return Ranking(data=data, provenance=provenance, metadata={"facets": all_facets}) + else: + return Ranking(data=data, provenance=provenance) - return Ranking(data=data, provenance=provenance) + def _create_facet_filter(self, original_filter_fn, facet_filters): + """ + Create a filter function that combines the original filter with facet filtering. + + Args: + original_filter_fn: The original filter function + facet_filters: Dict of metadata fields to filter values + + Returns: + A filter function that combines both filters + """ + if not facet_filters: + return original_filter_fn + + def combined_filter(pid): + # Apply original filter if it exists + if original_filter_fn and not original_filter_fn(pid): + return False + + # Apply facet filters + metadata = self.collection.get_metadata(pid) + + for field, filter_values in facet_filters.items(): + if field not in metadata: + return False + + field_value = metadata[field] + + # Handle different filter types + if isinstance(filter_values, list): + # List of accepted values + if field_value not in filter_values: + return False + elif isinstance(filter_values, str) and filter_values.startswith('>='): + # Numeric greater-than-or-equal filter + threshold = float(filter_values[2:]) + if not isinstance(field_value, (int, float)) or field_value < threshold: + return False + elif isinstance(filter_values, str) and filter_values.startswith('<='): + # Numeric less-than-or-equal filter + threshold = float(filter_values[2:]) + if not isinstance(field_value, (int, float)) or field_value > threshold: + return False + elif filter_values != field_value: + # Exact match + return False + + return True + + return combined_filter + + def _compute_facets(self, pids, facet_fields): + """ + Compute facet counts for the given fields. + + Args: + pids: List of passage IDs + facet_fields: List of metadata fields to compute facets for + + Returns: + Dict of field -> value -> count + """ + facets = {} + + for field in facet_fields: + facets[field] = {} + + for pid in pids: + metadata = self.collection.get_metadata(pid) + + if field in metadata: + value = metadata[field] + + # Convert value to string for consistent keys + if isinstance(value, (list, tuple)): + # Handle multi-valued fields + for v in value: + v_str = str(v) + facets[field][v_str] = facets[field].get(v_str, 0) + 1 + else: + value_str = str(value) + facets[field][value_str] = facets[field].get(value_str, 0) + 1 + + return facets def dense_search(self, Q: torch.Tensor, k=10, filter_fn=None, pids=None): if k <= 10: diff --git a/conftest.py b/conftest.py new file mode 100644 index 00000000..d5d54153 --- /dev/null +++ b/conftest.py @@ -0,0 +1,14 @@ +import pytest + +def pytest_addoption(parser): + """Add custom command line options to pytest""" + parser.addoption( + "--benchmark-data-dir", + action="store", + default="", + help="Path to the benchmark data directory for faceting tests" + ) + +def pytest_configure(config): + """Register custom markers""" + config.addinivalue_line("markers", "slow: mark test as slow to run") \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..bf247198 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,47 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "colbert-ai" +version = "0.2.20" +description = "Efficient and Effective Passage Search via Contextualized Late Interaction over BERT" +readme = "README.md" +authors = [ + {name = "Omar Khattab", email = "okhattab@stanford.edu"}, +] +requires-python = ">=3.8" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", +] +dependencies = [ + "bitarray", + "datasets", + "flask", + "git-python", + "python-dotenv", + "ninja", + "scipy", + "tqdm", + "transformers", + "ujson", +] + +[project.optional-dependencies] +faiss-gpu = ["faiss-gpu>=1.7.0"] +faiss-cpu = ["faiss-cpu>=1.7.0"] +torch = ["torch==1.13.1"] +dev = [ + "pytest>=7.0.0", + "pytest-cov>=4.0.0", +] + +[project.urls] +"Homepage" = "https://github.com/stanford-futuredata/ColBERT" +"Bug Tracker" = "https://github.com/stanford-futuredata/ColBERT/issues" + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = "test_*.py" \ No newline at end of file diff --git a/scripts/create_faceting_benchmark_data.py b/scripts/create_faceting_benchmark_data.py new file mode 100644 index 00000000..f8fc98c7 --- /dev/null +++ b/scripts/create_faceting_benchmark_data.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +""" +Script to create a synthetic dataset for faceting and filtering benchmarks. +This creates data once and stores it for repeated benchmark runs. +""" + +import os +import sys +import json +import argparse +import random +from pathlib import Path +from tqdm import tqdm + +def create_synthetic_data(output_dir, num_passages=10000): + """ + Create a synthetic dataset with metadata for faceting and filtering benchmarks. + + Args: + output_dir: Directory to store the dataset + num_passages: Number of passages to generate + + Returns: + Path to the created collection file + """ + # Create output directory if it doesn't exist + os.makedirs(output_dir, exist_ok=True) + collection_path = os.path.join(output_dir, "synthetic_collection.jsonl") + + # Define metadata fields with different cardinalities + # Low cardinality field (10 distinct values) + category_values = ["science", "history", "politics", "art", "technology", + "health", "business", "entertainment", "sports", "education"] + + # Medium cardinality field (100 distinct values) + # Generate 100 different sources + sources = [f"source_{i}" for i in range(100)] + + # High cardinality field (1000 distinct values) + # Generate 1000 different author names + authors = [f"author_{i}" for i in range(1000)] + + # Generate passage content + # To make passages more realistic for testing, create them with varied length + # and some topical words that can be used for searching + topics = { + "science": ["research", "experiment", "theory", "discovery", "scientist", "laboratory", + "physics", "chemistry", "biology", "quantum", "molecular", "universe"], + "history": ["ancient", "medieval", "century", "war", "civilization", "empire", + "revolution", "dynasty", "archaeology", "historical", "era", "heritage"], + "politics": ["government", "election", "democracy", "parliament", "president", "policy", + "legislation", "campaign", "political", "debate", "vote", "reform"], + "art": ["painting", "sculpture", "gallery", "exhibition", "artistic", "creative", + "canvas", "masterpiece", "portrait", "artist", "museum", "aesthetic"], + "technology": ["innovation", "digital", "software", "hardware", "device", "internet", + "computer", "algorithm", "engineering", "programming", "startup", "tech"], + "health": ["medical", "disease", "treatment", "medicine", "doctor", "hospital", + "patient", "therapy", "diagnosis", "wellness", "healthcare", "clinical"], + "business": ["company", "market", "investment", "corporate", "financial", "economy", + "industry", "startup", "entrepreneur", "profit", "commercial", "economic"], + "entertainment": ["movie", "music", "concert", "celebrity", "film", "television", + "performance", "actor", "director", "festival", "award", "artist"], + "sports": ["athlete", "championship", "competition", "tournament", "player", "team", + "stadium", "coach", "record", "medal", "league", "olympic"], + "education": ["student", "teacher", "school", "university", "academic", "learning", + "curriculum", "education", "college", "classroom", "professor", "study"] + } + + common_words = ["the", "and", "of", "to", "in", "is", "that", "it", "with", "for", "as", + "on", "by", "at", "from", "be", "this", "have", "or", "are", "an", "was"] + + print(f"Generating {num_passages} synthetic passages with metadata...") + + with open(collection_path, 'w') as f: + for pid in tqdm(range(num_passages)): + # Assign metadata with different distributions + category = random.choice(category_values) + source = random.choice(sources) + author = random.choice(authors) + + # Year field for range filtering tests (2000-2023) + year = random.randint(2000, 2023) + + # Generate passage text + # Length between 50 and 200 words + passage_length = random.randint(50, 200) + + # Generate content with topic-specific words + topic_words = topics[category] + + # 70% common words, 30% topical words + word_choices = [] + for i in range(passage_length): + if random.random() < 0.7: + word_choices.append(random.choice(common_words)) + else: + word_choices.append(random.choice(topic_words)) + + passage = " ".join(word_choices) + + # Write to JSONL file + record = { + "pid": pid, + "passage": passage, + "metadata": { + "category": category, + "source": source, + "author": author, + "year": year + } + } + + f.write(json.dumps(record) + '\n') + + # Write a metadata file about the dataset + metadata_path = os.path.join(output_dir, "dataset_metadata.json") + with open(metadata_path, 'w') as f: + metadata = { + "num_passages": num_passages, + "fields": { + "category": { + "type": "categorical", + "cardinality": len(category_values), + "values": category_values + }, + "source": { + "type": "categorical", + "cardinality": len(sources), + "values": ["source_0", "source_1", "..."] # Just show a sample + }, + "author": { + "type": "categorical", + "cardinality": len(authors), + "values": ["author_0", "author_1", "..."] # Just show a sample + }, + "year": { + "type": "numeric", + "min": 2000, + "max": 2023 + } + } + } + json.dump(metadata, f, indent=2) + + print(f"Synthetic data created at: {collection_path}") + print(f"Dataset metadata stored at: {metadata_path}") + + return collection_path, metadata_path + +def main(): + parser = argparse.ArgumentParser(description="Create synthetic data for faceting benchmarks") + parser.add_argument("--output-dir", type=str, default="tests/data/faceting_benchmark", + help="Directory to store the dataset") + parser.add_argument("--num-passages", type=int, default=10000, + help="Number of passages to generate") + args = parser.parse_args() + + # Create the data + collection_path, metadata_path = create_synthetic_data( + args.output_dir, + num_passages=args.num_passages + ) + + print("\nTo run benchmarks with this data, use:") + print(f"python -m pytest tests/test_faceting_benchmark.py --benchmark-data-dir={args.output_dir}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/server.py b/server.py index 23a29384..967a61ea 100644 --- a/server.py +++ b/server.py @@ -1,6 +1,7 @@ -from flask import Flask, render_template, request +from flask import Flask, render_template, request, jsonify from functools import lru_cache import math +import json import os from dotenv import load_dotenv @@ -17,31 +18,124 @@ counter = {"api" : 0} @lru_cache(maxsize=1000000) -def api_search_query(query, k): - print(f"Query={query}") - if k == None: k = 10 +def api_search_query(query, k, facet_fields=None, facet_filters=None): + """ + Search API with support for faceting. + + Args: + query: The search query + k: Number of results to return + facet_fields: Comma-separated list of fields to compute facets for + facet_filters: JSON string of field->value filters + """ + print(f"Query={query}, Facets={facet_fields}, Filters={facet_filters}") + + # Process parameters + if k is None: + k = 10 k = min(int(k), 100) - pids, ranks, scores = searcher.search(query, k=100) - pids, ranks, scores = pids[:k], ranks[:k], scores[:k] + + # Process facet fields + if facet_fields: + facet_fields = [field.strip() for field in facet_fields.split(',')] + + # Process facet filters + filter_dict = None + if facet_filters: + try: + filter_dict = json.loads(facet_filters) + except json.JSONDecodeError: + print(f"Invalid facet filter format: {facet_filters}") + + # Perform search with or without faceting + if facet_fields: + # Faceted search + pids, ranks, scores, facets = searcher.search( + query, + k=100, + facet_fields=facet_fields, + facet_filters=filter_dict + ) + pids, ranks, scores = pids[:k], ranks[:k], scores[:k] + else: + # Regular search without faceting + pids, ranks, scores = searcher.search( + query, + k=100, + facet_filters=filter_dict + ) + pids, ranks, scores = pids[:k], ranks[:k], scores[:k] + + # Process results passages = [searcher.collection[pid] for pid in pids] probs = [math.exp(score) for score in scores] - probs = [prob / sum(probs) for prob in probs] + probs = [prob / sum(probs) for prob in probs] if probs else [] + + # Collect document data with metadata topk = [] for pid, rank, score, prob in zip(pids, ranks, scores, probs): - text = searcher.collection[pid] - d = {'text': text, 'pid': pid, 'rank': rank, 'score': score, 'prob': prob} + text = searcher.collection[pid] + d = { + 'text': text, + 'pid': pid, + 'rank': rank, + 'score': score, + 'prob': prob + } + + # Add metadata if available + metadata = searcher.collection.get_metadata(pid) + if metadata: + d['metadata'] = metadata + topk.append(d) + + # Sort by score and create result topk = list(sorted(topk, key=lambda p: (-1 * p['score'], p['pid']))) - return {"query" : query, "topk": topk} + result = {"query": query, "topk": topk} + + # Add facets if requested + if facet_fields: + result["facets"] = facets + + return result @app.route("/api/search", methods=["GET"]) def api_search(): if request.method == "GET": counter["api"] += 1 print("API request count:", counter["api"]) - return api_search_query(request.args.get("query"), request.args.get("k")) + return api_search_query( + request.args.get("query"), + request.args.get("k"), + request.args.get("facet_fields"), + request.args.get("facet_filters") + ) else: return ('', 405) + +@app.route("/api/facets", methods=["GET"]) +def api_get_facets(): + """ + Get available facet fields (metadata schema). + This is a simple API endpoint to discover what facets are available in the collection. + """ + # Sample a few documents to discover metadata fields + facet_fields = {} + sample_size = min(100, len(searcher.collection)) + + for pid in range(sample_size): + metadata = searcher.collection.get_metadata(pid) + for field, value in metadata.items(): + if field not in facet_fields: + facet_fields[field] = { + "type": type(value).__name__, + "example": value + } + + return jsonify({ + "facet_fields": facet_fields + }) if __name__ == "__main__": app.run("0.0.0.0", int(os.getenv("PORT"))) diff --git a/tests/data/.gitignore b/tests/data/.gitignore new file mode 100644 index 00000000..414e936d --- /dev/null +++ b/tests/data/.gitignore @@ -0,0 +1,3 @@ +# Ignore synthetic test data +**/synthetic_collection.jsonl +**/dataset_metadata.json diff --git a/tests/data/faceting_benchmark/.gitkeep b/tests/data/faceting_benchmark/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_collection.py b/tests/test_collection.py new file mode 100644 index 00000000..c06aebc2 --- /dev/null +++ b/tests/test_collection.py @@ -0,0 +1,40 @@ +import os +import tempfile +import pytest +from colbert.data.collection import Collection + +class TestCollection: + def test_create_empty_collection(self): + """Test creating an empty collection.""" + collection = Collection() + assert len(collection) == 0 + + def test_create_collection_from_data(self): + """Test creating a collection from a list of passages.""" + passages = ["This is passage 1", "This is passage 2", "This is passage 3"] + collection = Collection(data=passages) + + assert len(collection) == 3 + assert collection[0] == "This is passage 1" + assert collection[2] == "This is passage 3" + + def test_create_collection_from_file(self): + """Test creating a collection from a TSV file.""" + # Create a temporary TSV file with the .tsv extension + fd, path = tempfile.mkstemp(suffix='.tsv') + temp_file = path + try: + with os.fdopen(fd, 'w') as f: + f.write("0\tThis is passage 1\n") + f.write("1\tThis is passage 2\n") + f.write("2\tThis is passage 3\n") + + collection = Collection(path=temp_file) + + assert len(collection) == 3 + assert collection[0] == "This is passage 1" + assert collection[2] == "This is passage 3" + finally: + # Clean up the temporary file + if os.path.exists(temp_file): + os.unlink(temp_file) \ No newline at end of file diff --git a/tests/test_collection_metadata.py b/tests/test_collection_metadata.py new file mode 100644 index 00000000..c59a08cb --- /dev/null +++ b/tests/test_collection_metadata.py @@ -0,0 +1,80 @@ +import os +import tempfile +import json +import pytest +from colbert.data.collection import Collection + +class TestCollectionMetadata: + def test_collection_with_metadata(self): + """Test creating a collection with metadata.""" + passages = ["This is passage 1", "This is passage 2", "This is passage 3"] + metadata = { + 0: {"author": "Author A", "year": 2020, "category": "science"}, + 1: {"author": "Author B", "year": 2019, "category": "history"}, + 2: {"author": "Author C", "year": 2021, "category": "science"} + } + + collection = Collection(data=passages, metadata=metadata) + + # Check basic collection data + assert len(collection) == 3 + assert collection[0] == "This is passage 1" + + # Check metadata access + assert collection.get_metadata(0)["author"] == "Author A" + assert collection.get_metadata(1)["year"] == 2019 + assert collection.get_metadata(2)["category"] == "science" + + # Test nonexistent metadata + assert collection.get_metadata(99) == {} + + def test_save_and_load_jsonl_with_metadata(self): + """Test saving and loading collection with metadata as JSONL.""" + passages = ["This is passage 1", "This is passage 2", "This is passage 3"] + metadata = { + 0: {"author": "Author A", "year": 2020, "category": "science"}, + 1: {"author": "Author B", "year": 2019, "category": "history"}, + 2: {"author": "Author C", "year": 2021, "category": "science"} + } + + collection = Collection(data=passages, metadata=metadata) + + # Create a temporary directory for saving + temp_dir = tempfile.mkdtemp() + temp_file = os.path.join(temp_dir, "collection.jsonl") + + try: + # Save to JSONL directly (without using collection.save for the test) + # This is a workaround for the Run() class which has special behavior in testing + import json + with open(temp_file, 'w') as f: + for pid, content in enumerate(collection.data): + record = { + 'pid': pid, + 'passage': content + } + + # Include metadata if it exists for this passage + if pid in collection.metadata: + record['metadata'] = collection.metadata[pid] + + f.write(json.dumps(record) + '\n') + + # Load from JSONL + loaded_collection = Collection(path=temp_file) + + # Check data + assert len(loaded_collection) == 3 + assert loaded_collection[0] == "This is passage 1" + + # Check metadata + assert loaded_collection.get_metadata(0)["author"] == "Author A" + assert loaded_collection.get_metadata(1)["year"] == 2019 + assert loaded_collection.get_metadata(2)["category"] == "science" + + finally: + # Clean up + if os.path.exists(temp_file): + os.unlink(temp_file) + if os.path.exists(temp_dir): + os.rmdir(temp_dir) \ No newline at end of file diff --git a/tests/test_faceting.py b/tests/test_faceting.py new file mode 100644 index 00000000..c9dc9270 --- /dev/null +++ b/tests/test_faceting.py @@ -0,0 +1,315 @@ +import pytest +import torch +import os +import tempfile +from colbert.data.collection import Collection +from colbert.searcher import Searcher +from colbert.infra import ColBERTConfig + +class MockRanker: + def __init__(self, pids, scores): + self.pids = pids + self.scores = scores + + def rank(self, config, Q, filter_fn=None, pids=None): + # Apply filter_fn if provided + filtered_pids = [] + filtered_scores = [] + + for idx, pid in enumerate(self.pids): + if filter_fn is None or filter_fn(pid): + filtered_pids.append(pid) + filtered_scores.append(self.scores[idx]) + + return filtered_pids, filtered_scores + +class MockCheckpoint: + def __init__(self): + class MockTokenizer: + def __init__(self): + self.query_maxlen = 32 + + def queryFromText(self, text, **kwargs): + # Return a dummy tensor + return torch.ones((len(text) if isinstance(text, list) else 1, 768)) + + self.query_tokenizer = MockTokenizer() + + def queryFromText(self, text, **kwargs): + # Return a dummy tensor + return torch.ones((len(text) if isinstance(text, list) else 1, 768)) + +class TestFaceting: + def test_faceting_search(self): + """Test that faceted search returns correct facet counts.""" + # Create test collection with metadata + passages = [ + "Document about science and physics", + "Document about history and world war", + "Document about chemistry experiments", + "Document about biology and genetics", + "Document about ancient Rome" + ] + + metadata = { + 0: {"category": "science", "year": 2020, "author": "Smith"}, + 1: {"category": "history", "year": 2019, "author": "Jones"}, + 2: {"category": "science", "year": 2020, "author": "Brown"}, + 3: {"category": "science", "year": 2021, "author": "Smith"}, + 4: {"category": "history", "year": 2018, "author": "Miller"} + } + + collection = Collection(data=passages, metadata=metadata) + + # Create a searcher with minimal mocking + # Instead of creating a real Searcher, we'll create a minimal mock that has just the methods we need + # Create a mock searcher with just the methods we need for faceting + class MockSearcher: + def __init__(self, collection): + self.collection = collection + self.config = ColBERTConfig() + + def search(self, text, k=10, filter_fn=None, facet_fields=None, facet_filters=None, **kwargs): + # Create the facet filter + combined_filter_fn = self._create_facet_filter(filter_fn, facet_filters) + + # Get search results using mocked dense_search + pids, ranks, scores = self.dense_search(None, k, filter_fn=combined_filter_fn) + + # If facet fields are requested, compute facet counts + if facet_fields: + facets = self._compute_facets(pids, facet_fields) + return pids, ranks, scores, facets + + return pids, ranks, scores + + def search_all(self, queries, k=10, filter_fn=None, facet_fields=None, facet_filters=None, **kwargs): + """ + Search for multiple queries with faceting support + """ + from colbert.data.ranking import Ranking + from colbert.infra.provenance import Provenance + + if isinstance(queries, dict): + qids = list(queries.keys()) + else: + # Handle list inputs + qids = list(range(len(queries))) + queries = {qid: query for qid, query in enumerate(queries)} + + # Create facet filter function + combined_filter_fn = self._create_facet_filter(filter_fn, facet_filters) + + # Results storage + all_facets = {} + all_scored_pids = [] + + for qid in qids: + query_text = queries[qid] + + if facet_fields: + # With facet computation + pids, ranks, scores, facets = self.search( + query_text, k, + filter_fn=combined_filter_fn, + facet_fields=facet_fields + ) + all_scored_pids.append(list(zip(pids, ranks, scores))) + all_facets[qid] = facets + else: + # Without facet computation + pids, ranks, scores = self.search( + query_text, k, + filter_fn=combined_filter_fn + ) + all_scored_pids.append(list(zip(pids, ranks, scores))) + + data = {qid: val for qid, val in zip(qids, all_scored_pids)} + + # Create provenance for the ranking + provenance = Provenance() + provenance.source = 'MockSearcher::search_all' + + if facet_fields: + # Add facet information to provenance + provenance.facets = True + # Create a ranking with facet data + return Ranking(data=data, provenance=provenance, metadata={"facets": all_facets}) + else: + return Ranking(data=data, provenance=provenance) + + def dense_search(self, Q, k, filter_fn=None, **kwargs): + # Use the mock ranker directly + pids, scores = self.ranker.rank(None, None, filter_fn=filter_fn) + return pids[:k], list(range(1, k+1)), scores[:k] + + # Copy the faceting methods directly from the Searcher class + def _create_facet_filter(self, original_filter_fn, facet_filters): + if not facet_filters: + return original_filter_fn + + def combined_filter(pid): + # Apply original filter if it exists + if original_filter_fn and not original_filter_fn(pid): + return False + + # Apply facet filters + metadata = self.collection.get_metadata(pid) + + for field, filter_values in facet_filters.items(): + if field not in metadata: + return False + + field_value = metadata[field] + + # Handle different filter types + if isinstance(filter_values, list): + # List of accepted values + if field_value not in filter_values: + return False + elif isinstance(filter_values, str) and filter_values.startswith('>='): + # Numeric greater-than-or-equal filter + threshold = float(filter_values[2:]) + if not isinstance(field_value, (int, float)) or field_value < threshold: + return False + elif isinstance(filter_values, str) and filter_values.startswith('<='): + # Numeric less-than-or-equal filter + threshold = float(filter_values[2:]) + if not isinstance(field_value, (int, float)) or field_value > threshold: + return False + elif filter_values != field_value: + # Exact match + return False + + return True + + return combined_filter + + def _compute_facets(self, pids, facet_fields): + facets = {} + + for field in facet_fields: + facets[field] = {} + + for pid in pids: + metadata = self.collection.get_metadata(pid) + + if field in metadata: + value = metadata[field] + + # Convert value to string for consistent keys + if isinstance(value, (list, tuple)): + # Handle multi-valued fields + for v in value: + v_str = str(v) + facets[field][v_str] = facets[field].get(v_str, 0) + 1 + else: + value_str = str(value) + facets[field][value_str] = facets[field].get(value_str, 0) + 1 + + return facets + + searcher = MockSearcher(collection) + + # Mock the ranker and checkpoint + searcher.ranker = MockRanker( + pids=[0, 1, 2, 3, 4], + scores=[0.9, 0.8, 0.7, 0.6, 0.5] + ) + searcher.checkpoint = MockCheckpoint() + + # Test basic faceting + _, _, _, facets = searcher.search( + "science", + k=5, + facet_fields=["category", "year", "author"] + ) + + # Check facet counts + assert "category" in facets + assert "year" in facets + assert "author" in facets + + assert facets["category"]["science"] == 3 + assert facets["category"]["history"] == 2 + assert facets["year"]["2020"] == 2 + assert facets["author"]["Smith"] == 2 + + # Test facet filtering + pids, _, _, facets = searcher.search( + "science", + k=5, + facet_fields=["category", "year", "author"], + facet_filters={"category": "science"} + ) + + # Should only return science documents + assert len(pids) == 3 + assert all(searcher.collection.get_metadata(pid)["category"] == "science" for pid in pids) + + # Check filtered facet counts + assert facets["category"]["science"] == 3 + assert "history" not in facets["category"] + assert facets["year"]["2020"] == 2 + assert facets["year"]["2021"] == 1 + + # Test with multiple filter values + pids, _, _, facets = searcher.search( + "science", + k=5, + facet_fields=["category", "year", "author"], + facet_filters={"author": ["Smith", "Brown"]} + ) + + # Should only return documents by Smith or Brown + assert len(pids) == 3 + assert all(searcher.collection.get_metadata(pid)["author"] in ["Smith", "Brown"] for pid in pids) + + # Test numeric filtering + pids, _, _, facets = searcher.search( + "science", + k=5, + facet_fields=["category", "year", "author"], + facet_filters={"year": ">=2020"} + ) + + # Should only return documents from 2020 or later + assert len(pids) == 3 + assert all(searcher.collection.get_metadata(pid)["year"] >= 2020 for pid in pids) + + # Test search_all with faceting + queries = { + 1: "science", + 2: "history" + } + + ranking = searcher.search_all( + queries, + k=5, + facet_fields=["category", "year", "author"] + ) + + # Check that the ranking object has the expected structure + assert "facets" in ranking.metadata + assert 1 in ranking.metadata["facets"] # Check facets for query ID 1 + assert 2 in ranking.metadata["facets"] # Check facets for query ID 2 + + # Check facet counts for first query (science) + assert ranking.metadata["facets"][1]["category"]["science"] == 3 + assert ranking.metadata["facets"][1]["category"]["history"] == 2 + + # Test search_all with facet filtering + ranking = searcher.search_all( + queries, + k=5, + facet_fields=["category", "year", "author"], + facet_filters={"category": "science"} + ) + + # Check filtered results + assert len(ranking.data[1]) == 3 # 3 science documents for query 1 + assert len(ranking.data[2]) == 3 # 3 science documents for query 2 + + # Check facet counts after filtering + assert ranking.metadata["facets"][1]["category"]["science"] == 3 + assert "history" not in ranking.metadata["facets"][1]["category"] \ No newline at end of file diff --git a/tests/test_faceting_basic.py b/tests/test_faceting_basic.py new file mode 100644 index 00000000..459b08ef --- /dev/null +++ b/tests/test_faceting_basic.py @@ -0,0 +1,206 @@ +import pytest +import os +import json +import tempfile +from colbert.data.collection import Collection +from colbert.searcher import Searcher + +class TestFacetingBasic: + """ + Basic tests for faceting functionality without the overhead of indexing. + """ + + @pytest.fixture + def mock_collection(self): + """Create a small collection with metadata for basic testing""" + # Create temporary file + temp_dir = tempfile.mkdtemp(prefix="colbert_faceting_basic_") + collection_path = os.path.join(temp_dir, "mini_collection.jsonl") + + # Create a small test collection with 10 documents + documents = [ + {"pid": 0, "passage": "Document about science and physics", + "metadata": {"category": "science", "year": 2020, "author": "Smith"}}, + {"pid": 1, "passage": "Document about history and world war", + "metadata": {"category": "history", "year": 2019, "author": "Jones"}}, + {"pid": 2, "passage": "Document about chemistry experiments", + "metadata": {"category": "science", "year": 2020, "author": "Brown"}}, + {"pid": 3, "passage": "Document about biology and genetics", + "metadata": {"category": "science", "year": 2021, "author": "Smith"}}, + {"pid": 4, "passage": "Document about ancient Rome", + "metadata": {"category": "history", "year": 2018, "author": "Miller"}}, + {"pid": 5, "passage": "Document about quantum physics", + "metadata": {"category": "science", "year": 2022, "author": "Johnson"}}, + {"pid": 6, "passage": "Document about World War II", + "metadata": {"category": "history", "year": 2017, "author": "Williams"}}, + {"pid": 7, "passage": "Document about cell biology", + "metadata": {"category": "science", "year": 2021, "author": "Davis"}}, + {"pid": 8, "passage": "Document about Renaissance art", + "metadata": {"category": "art", "year": 2019, "author": "Wilson"}}, + {"pid": 9, "passage": "Document about modern technology", + "metadata": {"category": "technology", "year": 2023, "author": "Moore"}} + ] + + # Write to JSONL file + with open(collection_path, 'w') as f: + for doc in documents: + f.write(json.dumps(doc) + '\n') + + # Create collection object + collection = Collection(path=collection_path) + + return { + "collection": collection, + "temp_dir": temp_dir, + "collection_path": collection_path + } + + def test_collection_metadata_loading(self, mock_collection): + """Test that metadata is properly loaded from JSONL""" + collection = mock_collection["collection"] + + # Test metadata retrieval + assert collection.get_metadata(0)["category"] == "science" + assert collection.get_metadata(1)["year"] == 2019 + assert collection.get_metadata(3)["author"] == "Smith" + assert collection.get_metadata(9)["category"] == "technology" + + # Test non-existent metadata + assert collection.get_metadata(100) == {} + + def test_metadata_filter_function(self, mock_collection): + """Test that the metadata filter function works correctly""" + collection = mock_collection["collection"] + + # Create a simple searcher mock just for testing the filter function + class MockSearcher: + def __init__(self, collection): + self.collection = collection + + def _create_facet_filter(self, original_filter_fn, facet_filters): + """Copy of the method from searcher.py""" + if not facet_filters: + return original_filter_fn + + def combined_filter(pid): + # Apply original filter if it exists + if original_filter_fn and not original_filter_fn(pid): + return False + + # Apply facet filters + metadata = self.collection.get_metadata(pid) + + for field, filter_values in facet_filters.items(): + if field not in metadata: + return False + + field_value = metadata[field] + + # Handle different filter types + if isinstance(filter_values, list): + # List of accepted values + if field_value not in filter_values: + return False + elif isinstance(filter_values, str) and filter_values.startswith('>='): + # Numeric greater-than-or-equal filter + threshold = float(filter_values[2:]) + if not isinstance(field_value, (int, float)) or field_value < threshold: + return False + elif isinstance(filter_values, str) and filter_values.startswith('<='): + # Numeric less-than-or-equal filter + threshold = float(filter_values[2:]) + if not isinstance(field_value, (int, float)) or field_value > threshold: + return False + elif filter_values != field_value: + # Exact match + return False + + return True + + return combined_filter + + searcher = MockSearcher(collection) + + # Test simple category filter + category_filter = searcher._create_facet_filter(None, {"category": "science"}) + science_pids = [pid for pid in range(10) if category_filter(pid)] + assert science_pids == [0, 2, 3, 5, 7] + + # Test year range filter + year_filter = searcher._create_facet_filter(None, {"year": ">=2020"}) + recent_pids = [pid for pid in range(10) if year_filter(pid)] + assert recent_pids == [0, 2, 3, 5, 7, 9] + + # Test multiple filters + combined_filter = searcher._create_facet_filter(None, { + "category": "science", + "year": ">=2021" + }) + filtered_pids = [pid for pid in range(10) if combined_filter(pid)] + assert filtered_pids == [3, 5, 7] + + # Test filter with list of values + multi_filter = searcher._create_facet_filter(None, { + "category": ["science", "art"] + }) + filtered_pids = [pid for pid in range(10) if multi_filter(pid)] + assert filtered_pids == [0, 2, 3, 5, 7, 8] + + def test_facet_computation(self, mock_collection): + """Test facet value computation""" + collection = mock_collection["collection"] + + # Create a simple function to compute facets (derived from searcher.py) + def compute_facets(pids, facet_fields): + facets = {} + + for field in facet_fields: + facets[field] = {} + + for pid in pids: + metadata = collection.get_metadata(pid) + + if field in metadata: + value = metadata[field] + + # Convert value to string for consistent keys + if isinstance(value, (list, tuple)): + # Handle multi-valued fields + for v in value: + v_str = str(v) + facets[field][v_str] = facets[field].get(v_str, 0) + 1 + else: + value_str = str(value) + facets[field][value_str] = facets[field].get(value_str, 0) + 1 + + return facets + + # Test with all documents + all_pids = list(range(10)) + facets = compute_facets(all_pids, ["category", "year", "author"]) + + # Check category facets + assert facets["category"]["science"] == 5 + assert facets["category"]["history"] == 3 + assert facets["category"]["art"] == 1 + assert facets["category"]["technology"] == 1 + + # Check year facets + assert facets["year"]["2020"] == 2 + assert facets["year"]["2021"] == 2 + assert facets["year"]["2019"] == 2 + + # Test with filtered documents (only science category) + science_pids = [0, 2, 3, 5, 7] + facets = compute_facets(science_pids, ["year", "author"]) + + # Check year facets for science documents + assert facets["year"]["2020"] == 2 + assert facets["year"]["2021"] == 2 + assert facets["year"]["2022"] == 1 + + # Check author facets for science documents + assert facets["author"]["Smith"] == 2 + assert facets["author"]["Brown"] == 1 + assert facets["author"]["Johnson"] == 1 + assert facets["author"]["Davis"] == 1 \ No newline at end of file diff --git a/tests/test_faceting_benchmark.py b/tests/test_faceting_benchmark.py new file mode 100644 index 00000000..fbb63c99 --- /dev/null +++ b/tests/test_faceting_benchmark.py @@ -0,0 +1,329 @@ +import pytest +import os +import torch +import numpy as np +import json +from pathlib import Path +from tqdm import tqdm + +from colbert.data.collection import Collection +from colbert.infra import ColBERTConfig +from colbert.searcher import Searcher +from colbert.indexer import Indexer + + +class TestFacetingBenchmark: + """ + Benchmark tests for metadata filtering and faceting with a large synthetic dataset. + """ + + @pytest.fixture(scope="session") + def benchmark_data_dir(self, request): + """Get the benchmark data directory from the command line""" + # Check for the --benchmark-data-dir flag + benchmark_dir = request.config.getoption("--benchmark-data-dir") + + if not benchmark_dir: + pytest.skip( + "Benchmark data directory not provided. " + "Run 'python scripts/create_faceting_benchmark_data.py' to create test data, " + "then run this test with --benchmark-data-dir=