Download and load BEIR dataset
(dataset: str)
| 469 | return TextEmbedding(model_name="thenlper/gte-base", max_length=512) |
| 470 | |
| 471 | def get_beir_dataset(dataset: str) -> Tuple[Dict, Dict, Dict]: |
| 472 | """Download and load BEIR dataset""" |
| 473 | save_dir = Path("datasets") / f"hybrid_{dataset}" |
| 474 | archive_dir = save_dir / "archive" |
| 475 | save_dir.mkdir(parents=True, exist_ok=True) |
| 476 | archive_dir.mkdir(parents=True, exist_ok=True) |
| 477 | |
| 478 | # Download dataset and unzip the dataset |
| 479 | base_url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip" |
| 480 | data_path = beir.util.download_and_unzip(base_url.format(dataset), str(archive_dir)) |
| 481 | |
| 482 | if dataset == "msmarco": |
| 483 | split = "dev" |
| 484 | else: |
| 485 | split = "test" |
| 486 | |
| 487 | # Merge cqadupstack subdatasets if needed |
| 488 | if dataset == "cqadupstack": |
| 489 | merge_cqa_dupstack(data_path, verbose=True) |
| 490 | |
| 491 | loader = GenericDataLoader(data_folder=data_path) |
| 492 | loader.check(fIn=loader.corpus_file, ext="jsonl") |
| 493 | loader.check(fIn=loader.query_file, ext="jsonl") |
| 494 | loader._load_corpus() |
| 495 | loader._load_queries() |
| 496 | corpus = loader.corpus |
| 497 | queries = loader.queries |
| 498 | |
| 499 | loader.qrels_file = os.path.join(loader.qrels_folder, split + ".tsv") |
| 500 | loader._load_qrels() |
| 501 | qrels = loader.qrels |
| 502 | |
| 503 | return corpus, queries, qrels |
| 504 | |
| 505 | def generate_dense_embeddings(corpus: Dict, queries: Dict, dataset: str, embedding_model) -> Tuple[List[str], np.ndarray, List[str], np.ndarray]: |
| 506 | """Generate dense embeddings for corpus and queries""" |
no test coverage detected