Function get_dataset

tests/test-tf-idf-bm25.py:187–230

Generate raw dataset with just term frequencies

(dataset: str)

def get_dataset(dataset: str) -> tuple:
    """Generate raw dataset with just term frequencies"""
    dataset_folder = f"datasets/sparse_idf_dataset_{dataset}"
    raw_dataset_file = f"{dataset_folder}/raw_vectors.pkl"
    if os.path.exists(raw_dataset_file):
        print(f"Loading existing raw dataset from {raw_dataset_file}")
        with open(raw_dataset_file, "rb") as f:
            return pickle.load(f)

    save_dir = f"{dataset_folder}/dataset"

    print("Downloading the dataset...")
    bm25s.utils.beir.download_dataset(dataset, save_dir=save_dir)
    print("Loading the corpus...")
    corpus = bm25s.utils.beir.load_corpus(dataset, save_dir=save_dir)

    punctuations = get_all_punctuation()
    stemmer = SnowballStemmer("english")

    print("Converting to sparse vectors...")
    vectors = []

    with ThreadPoolExecutor(max_workers=32) as executor:
        futures = []
        for k, v in tqdm(corpus.items()):
            futures.append(
                executor.submit(
                    transform_sentence_to_vector,
                    k,
                    v["title"] + " " + v["text"],
                    punctuations,
                    stemmer,
                )
            )

        for future in tqdm(as_completed(futures)):
            vectors.append(future.result())

    # Save to disk
    with open(raw_dataset_file, "wb") as f:
        pickle.dump(vectors, f)

    print(f"Raw dataset generated and saved to {raw_dataset_file}")
    return vectors


def calculate_corpus_statistics(dataset, vectors):