miroir/docs/research/score-comparability/queries/generate.py

#!/usr/bin/env python3
"""
Generate query sets for score comparability experiments.

Query types:
1. Single-term queries - test basic term frequency handling
2. Multi-term AND queries - test phrase matching
3. Category-filtered queries - test filtered search
4. Rare-term queries - test IDF behavior on low-frequency terms
5. Common-term queries - test IDF behavior on high-frequency terms
"""

import argparse
import json
import random
from pathlib import Path
from typing import List, Dict, Set


def load_vocabulary(corpus_dir: Path) -> Dict:
    """Load vocabulary and corpus metadata."""
    vocab_file = corpus_dir / "vocabulary.json"
    metadata_file = corpus_dir / "metadata.json"

    with open(vocab_file) as f:
        vocab_data = json.load(f)

    with open(metadata_file) as f:
        metadata = json.load(f)

    # Load corpus to compute term frequencies
    corpus_file = corpus_dir / "corpus.jsonl"
    term_freq = {}

    with open(corpus_file) as f:
        for line in f:
            doc = json.loads(line)
            text = f"{doc['title']} {doc['content']}".lower()
            words = set(text.split())
            for word in words:
                if word in vocab_data["terms"]:
                    term_freq[word] = term_freq.get(word, 0) + 1

    # Sort terms by frequency
    sorted_terms = sorted(term_freq.items(), key=lambda x: x[1])

    return {
        "terms": vocab_data["terms"],
        "categories": vocab_data["categories"],
        "term_freq": dict(sorted_terms),
        "total_docs": metadata["total_documents"],
    }


def generate_single_term_queries(vocab_data: Dict, count: int) -> List[Dict]:
    """Generate single-term queries with random term selection."""
    queries = []
    terms = vocab_data["terms"]

    for i in range(count):
        term = random.choice(terms)
        queries.append({
            "id": f"q-single-{i:05d}",
            "type": "single_term",
            "q": term,
            "filter": None,
        })

    return queries


def generate_multi_term_queries(vocab_data: Dict, count: int, min_terms: int = 2, max_terms: int = 4) -> List[Dict]:
    """Generate multi-term queries."""
    queries = []
    terms = vocab_data["terms"]

    for i in range(count):
        num_terms = random.randint(min_terms, max_terms)
        selected = random.sample(terms, min(num_terms, len(terms)))
        queries.append({
            "id": f"q-multi-{i:05d}",
            "type": "multi_term",
            "q": " ".join(selected),
            "filter": None,
        })

    return queries


def generate_filtered_queries(vocab_data: Dict, count: int) -> List[Dict]:
    """Generate queries with category filters."""
    queries = []
    terms = vocab_data["terms"]
    categories = vocab_data["categories"]

    for i in range(count):
        term = random.choice(terms)
        category = random.choice(categories)
        queries.append({
            "id": f"q-filter-{i:05d}",
            "type": "filtered",
            "q": term,
            "filter": f"category = {category}",
        })

    return queries


def generate_rare_term_queries(vocab_data: Dict, count: int, percentile: float = 0.1) -> List[Dict]:
    """Generate queries using rare terms (low document frequency)."""
    queries = []
    term_freq = vocab_data["term_freq"]
    sorted_terms = list(term_freq.items())

    # Get rare terms (bottom percentile by frequency)
    cutoff = int(len(sorted_terms) * percentile)
    rare_terms = [t for t, _ in sorted_terms[:cutoff]]

    for i in range(count):
        if not rare_terms:
            break
        term = random.choice(rare_terms)
        queries.append({
            "id": f"q-rare-{i:05d}",
            "type": "rare_term",
            "q": term,
            "filter": None,
        })

    return queries


def generate_common_term_queries(vocab_data: Dict, count: int, percentile: float = 0.9) -> List[Dict]:
    """Generate queries using common terms (high document frequency)."""
    queries = []
    term_freq = vocab_data["term_freq"]
    sorted_terms = list(term_freq.items())

    # Get common terms (top percentile by frequency)
    cutoff = int(len(sorted_terms) * percentile)
    common_terms = [t for t, _ in sorted_terms[cutoff:]]

    for i in range(count):
        if not common_terms:
            break
        term = random.choice(common_terms)
        queries.append({
            "id": f"q-common-{i:05d}",
            "type": "common_term",
            "q": term,
            "filter": None,
        })

    return queries


def main():
    parser = argparse.ArgumentParser(description="Generate query sets for experiments")
    parser.add_argument("--corpus", type=str, default="corpus/", help="Corpus directory")
    parser.add_argument("--output", type=str, default="queries/", help="Output directory")
    parser.add_argument("--total", type=int, default=10000, help="Total number of queries")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")

    args = parser.parse_args()

    random.seed(args.seed)

    corpus_dir = Path(args.corpus)
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Loading vocabulary from {corpus_dir}...")
    vocab_data = load_vocabulary(corpus_dir)

    print(f"Vocabulary: {len(vocab_data['terms'])} terms")
    print(f"Categories: {vocab_data['categories']}")
    print(f"Term frequency range: {min(vocab_data['term_freq'].values())} - {max(vocab_data['term_freq'].values())}")

    # Generate different query types
    print(f"\nGenerating {args.total} queries...")

    allocation = {
        "single_term": 0.25,
        "multi_term": 0.25,
        "filtered": 0.20,
        "rare_term": 0.15,
        "common_term": 0.15,
    }

    queries = []

    # Single-term queries
    count = int(args.total * allocation["single_term"])
    queries.extend(generate_single_term_queries(vocab_data, count))
    print(f"  Single-term: {count}")

    # Multi-term queries
    count = int(args.total * allocation["multi_term"])
    queries.extend(generate_multi_term_queries(vocab_data, count))
    print(f"  Multi-term: {count}")

    # Filtered queries
    count = int(args.total * allocation["filtered"])
    queries.extend(generate_filtered_queries(vocab_data, count))
    print(f"  Filtered: {count}")

    # Rare-term queries
    count = int(args.total * allocation["rare_term"])
    rare_queries = generate_rare_term_queries(vocab_data, count)
    queries.extend(rare_queries)
    print(f"  Rare-term: {len(rare_queries)}")

    # Common-term queries
    count = int(args.total * allocation["common_term"])
    common_queries = generate_common_term_queries(vocab_data, count)
    queries.extend(common_queries)
    print(f"  Common-term: {len(common_queries)}")

    # Shuffle to mix query types
    random.shuffle(queries)

    # Save query set
    output_file = output_dir / "queries.jsonl"
    with open(output_file, "w") as f:
        for q in queries:
            f.write(json.dumps(q) + "\n")

    # Save metadata
    metadata = {
        "total_queries": len(queries),
        "allocation": allocation,
        "random_seed": args.seed,
        "vocab_size": len(vocab_data["terms"]),
        "categories": vocab_data["categories"],
    }
    with open(output_dir / "metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    print(f"\nGenerated {len(queries)} queries")
    print(f"Saved to {output_file}")


if __name__ == "__main__":
    main()