Function load_queries_robust

tests/modal_remote_vllm/remote_cache_dense.py:128–199

Robust query loading that handles malformed JSON escape sequences. Fallback for when beir.load_queries fails due to JSON parsing errors.

(dataset: str, save_dir: str)

def load_queries_robust(dataset: str, save_dir: str) -> dict:
    """
    Robust query loading that handles malformed JSON escape sequences.
    Fallback for when beir.load_queries fails due to JSON parsing errors.
    """
    import json
    from pathlib import Path

    queries_file = Path(save_dir) / dataset / "queries.jsonl"
    if not queries_file.exists():
        raise FileNotFoundError(f"Queries file not found: {queries_file}")

    queries = {}
    skipped_lines = 0

    with open(queries_file, "r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue

            try:
                # Try standard JSON parsing first
                query_data = json.loads(line)
                query_id = query_data.get("_id", str(line_num))

                # Extract text, handling different possible field names
                text = (
                    query_data.get("text")
                    or query_data.get("query")
                    or query_data.get("body", "")
                )

                queries[query_id] = {"text": text}

            except json.JSONDecodeError:
                # Handle malformed escape sequences
                try:
                    # Fix common escape sequence issues
                    fixed_line = line
                    # Replace problematic \x sequences with unicode equivalents
                    import re

                    # Replace \xef with proper unicode
                    fixed_line = re.sub(
                        r"\\x([0-9a-fA-F]{2})",
                        lambda m: chr(int(m.group(1), 16)),
                        fixed_line,
                    )

                    query_data = json.loads(fixed_line)
                    query_id = query_data.get("_id", str(line_num))

                    text = (
                        query_data.get("text")
                        or query_data.get("query")
                        or query_data.get("body", "")
                    )