(
ds: str, embedder: RemoteEmbeddings, use_chunking: bool = False
)
| 305 | |
| 306 | |
| 307 | def cache_dataset( |
| 308 | ds: str, embedder: RemoteEmbeddings, use_chunking: bool = False |
| 309 | ) -> None: |
| 310 | cache_file = Path("datasets") / f"hybrid_{ds}" / "dense_embeddings.npz" |
| 311 | if cache_file.exists(): |
| 312 | print(f"[{ds}] already cached – skipping") |
| 313 | return |
| 314 | |
| 315 | corpus = ensure_corpus(ds) |
| 316 | |
| 317 | # Load queries as well |
| 318 | print(f"[{ds}] loading queries ...") |
| 319 | import bm25s.utils.beir as beir |
| 320 | |
| 321 | try: |
| 322 | queries = beir.load_queries(ds, save_dir=str(Path("datasets") / f"hybrid_{ds}")) |
| 323 | print(f"[{ds}] loaded {len(queries):,} queries") |
| 324 | except Exception as e: |
| 325 | print(f"[{ds}] Error loading queries with beir, attempting manual parsing: {e}") |
| 326 | queries = load_queries_robust( |
| 327 | ds, save_dir=str(Path("datasets") / f"hybrid_{ds}") |
| 328 | ) |
| 329 | print(f"[{ds}] loaded {len(queries):,} queries (robust parsing)") |
| 330 | |
| 331 | if use_chunking: |
| 332 | # Use chunking approach for very long documents |
| 333 | print(f"[{ds}] processing documents with chunking...") |
| 334 | texts = [] |
| 335 | corpus_ids = [] |
| 336 | for doc_id, doc in tqdm.tqdm( |
| 337 | corpus.items(), desc=f"{ds} chunking", unit="docs" |
| 338 | ): |
| 339 | combined_text = f"{doc.get('title', '')} {doc['text']}".strip() |
| 340 | chunks = chunk_text_by_tokens(combined_text, max_tokens=MAX_TOKENS) |
| 341 | |
| 342 | for i, chunk in enumerate(chunks): |
| 343 | texts.append(chunk) |
| 344 | # Create unique IDs for chunks |
| 345 | chunk_id = f"{doc_id}_chunk_{i}" if len(chunks) > 1 else doc_id |
| 346 | corpus_ids.append(chunk_id) |
| 347 | print( |
| 348 | f"[{ds}] created {len(texts):,} text chunks from {len(corpus):,} documents" |
| 349 | ) |
| 350 | else: |
| 351 | # Use simple clipping approach |
| 352 | print(f"[{ds}] processing documents with clipping...") |
| 353 | texts = [] |
| 354 | corpus_ids = list(corpus.keys()) |
| 355 | |
| 356 | for doc_id in tqdm.tqdm(corpus_ids, desc=f"{ds} clipping", unit="docs"): |
| 357 | doc = corpus[doc_id] |
| 358 | combined_text = f"{doc.get('title', '')} {doc['text']}".strip() |
| 359 | clipped_text = clip_to_max_tokens(combined_text, max_len=MAX_TOKENS) |
| 360 | |
| 361 | # Double-check the clipped text |
| 362 | if not validate_token_length(clipped_text, MAX_TOKENS): |
| 363 | print( |
| 364 | f"\n[{ds}] Warning: Document {doc_id} still too long, applying aggressive clipping" |
no test coverage detected