Add faceting and metadata support to ColBERT

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

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

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

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)