pyserini/util.py

#
# Pyserini: Reproducible IR research with sparse and dense representations
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import hashlib
import os
import re
import shutil
import tarfile
import logging
from urllib.error import HTTPError, URLError
from urllib.request import urlretrieve

import pandas as pd
from tqdm import tqdm

from pyserini.encoded_query_info import QUERY_INFO
from pyserini.encoded_corpus_info import CORPUS_INFO
from pyserini.evaluate_script_info import EVALUATION_INFO
from pyserini.prebuilt_index_info import TF_INDEX_INFO, FAISS_INDEX_INFO, IMPACT_INDEX_INFO


logger = logging.getLogger(__name__)


# https://gist.github.com/leimao/37ff6e990b3226c2c9670a2cd1e4a6f5
class TqdmUpTo(tqdm):
    def update_to(self, b=1, bsize=1, tsize=None):
        """
        b  : int, optional
            Number of blocks transferred so far [default: 1].
        bsize  : int, optional
            Size of each block (in tqdm units) [default: 1].
        tsize  : int, optional
            Total size (in tqdm units). If [default: None] remains unchanged.
        """
        if tsize is not None:
            self.total = tsize
        self.update(b * bsize - self.n)  # will also set self.n = b * bsize


# For large files, we need to compute MD5 block by block. See:
# https://stackoverflow.com/questions/1131220/get-md5-hash-of-big-files-in-python
def compute_md5(file, block_size=2**20):
    m = hashlib.md5()
    with open(file, 'rb') as f:
        while True:
            buf = f.read(block_size)
            if not buf:
                break
            m.update(buf)
    return m.hexdigest()


def download_url(url, save_dir, local_filename=None, md5=None, force=False, verbose=True):
    # If caller does not specify local filename, figure it out from the download URL:
    if not local_filename:
        filename = url.split('/')[-1]
        filename = re.sub('\\?dl=1$', '', filename)  # Remove the Dropbox 'force download' parameter
    else:
        # Otherwise, use the specified local_filename:
        filename = local_filename

    destination_path = os.path.join(save_dir, filename)

    if verbose:
        print(f'Downloading {url} to {destination_path}...')

    # Check to see if file already exists, if so, simply return (quietly) unless force=True, in which case we remove
    # destination file and download fresh copy.
    if os.path.exists(destination_path):
        if verbose:
            print(f'{destination_path} already exists!')
        if not force:
            if verbose:
                print(f'Skipping download.')
            return destination_path
        if verbose:
            print(f'force=True, removing {destination_path}; fetching fresh copy...')
        os.remove(destination_path)

    with TqdmUpTo(unit='B', unit_scale=True, unit_divisor=1024, miniters=1, desc=filename) as t:
        urlretrieve(url, filename=destination_path, reporthook=t.update_to)

    if md5:
        md5_computed = compute_md5(destination_path)
        assert md5_computed == md5, f'{destination_path} does not match checksum! Expecting {md5} got {md5_computed}.'

    return destination_path


def get_cache_home():
    custom_dir = os.environ.get("PYSERINI_CACHE")
    if custom_dir is not None and custom_dir != '':
        return custom_dir
    return os.path.expanduser(os.path.join(f'~{os.path.sep}.cache', "pyserini"))

def download_and_unpack_index(url, index_directory='indexes', local_filename=False,
                              force=False, verbose=True, prebuilt=False, md5=None):
    # If caller does not specify local filename, figure it out from the download URL:
    if not local_filename:
        index_name = url.split('/')[-1]
    else:
        # Otherwise, use the specified local_filename:
        index_name = local_filename
    # Remove the suffix:
    index_name = re.sub('''.tar.gz.*$''', '', index_name)

    if prebuilt:
        index_directory = os.path.join(get_cache_home(), index_directory)
        index_path = os.path.join(index_directory, f'{index_name}.{md5}')

        if not os.path.exists(index_directory):
            os.makedirs(index_directory)

        local_tarball = os.path.join(index_directory, f'{index_name}.tar.gz')
        # If there's a local tarball, it's likely corrupted, because we remove the local tarball on success (below).
        # So, we want to remove.
        if os.path.exists(local_tarball):
            os.remove(local_tarball)
    else:
        local_tarball = os.path.join(index_directory, f'{index_name}.tar.gz')
        index_path = os.path.join(index_directory, f'{index_name}')

    # Check to see if index already exists, if so, simply return (quietly) unless force=True, in which case we remove
    # index and download fresh copy.
    if os.path.exists(index_path):
        if not force:
            if verbose:
                print(f'{index_path} already exists, skipping download.')
            return index_path
        if verbose:
            print(f'{index_path} already exists, but force=True, removing {index_path} and fetching fresh copy...')
        shutil.rmtree(index_path)

    print(f'Downloading index at {url}...')
    download_url(url, index_directory, local_filename=local_filename, verbose=False, md5=md5)

    if verbose:
        print(f'Extracting {local_tarball} into {index_path}...')
    try:
        tarball = tarfile.open(local_tarball)
    except:
        local_tarball = os.path.join(index_directory, f'{index_name}')
        tarball = tarfile.open(local_tarball)
    dirs_in_tarball = [member.name for member in tarball if member.isdir()]
    assert len(dirs_in_tarball), f"Detect multiple members ({', '.join(dirs_in_tarball)}) under the tarball {local_tarball}."
    tarball.extractall(index_directory)
    tarball.close()
    os.remove(local_tarball)

    if prebuilt:
        dir_in_tarball = dirs_in_tarball[0]
        if dir_in_tarball != index_name:
            logger.info(f"Renaming {index_directory}/{dir_in_tarball} into {index_directory}/{index_name}.")
            index_name = dir_in_tarball
        os.rename(os.path.join(index_directory, f'{index_name}'), index_path)

    return index_path


def check_downloaded(index_name):
    if index_name in TF_INDEX_INFO:
        target_index = TF_INDEX_INFO[index_name]
    elif index_name in IMPACT_INDEX_INFO:
        target_index = IMPACT_INDEX_INFO[index_name]
    else:
        target_index = FAISS_INDEX_INFO[index_name]
    index_url = target_index['urls'][0]
    index_md5 = target_index['md5']
    index_name = index_url.split('/')[-1]
    index_name = re.sub('''.tar.gz.*$''', '', index_name)
    index_directory = os.path.join(get_cache_home(), 'indexes')
    index_path = os.path.join(index_directory, f'{index_name}.{index_md5}')

    return os.path.exists(index_path)


def get_sparse_indexes_info():
    df = pd.DataFrame.from_dict({**TF_INDEX_INFO, **IMPACT_INDEX_INFO})
    for index in df.keys():
        df[index]['downloaded'] = check_downloaded(index)

    with pd.option_context('display.max_rows', None, 'display.max_columns',
                           None, 'display.max_colwidth', None, 'display.colheader_justify', 'left'):
        print(df)


def get_impact_indexes_info():
    df = pd.DataFrame.from_dict(IMPACT_INDEX_INFO)
    for index in df.keys():
        df[index]['downloaded'] = check_downloaded(index)

    with pd.option_context('display.max_rows', None, 'display.max_columns',
                           None, 'display.max_colwidth', None, 'display.colheader_justify', 'left'):
        print(df)


def get_dense_indexes_info():
    df = pd.DataFrame.from_dict(FAISS_INDEX_INFO)
    for index in df.keys():
        df[index]['downloaded'] = check_downloaded(index)

    with pd.option_context('display.max_rows', None, 'display.max_columns',
                           None, 'display.max_colwidth', None, 'display.colheader_justify', 'left'):
        print(df)


def download_prebuilt_index(index_name, force=False, verbose=True, mirror=None):
    if index_name not in TF_INDEX_INFO and index_name not in FAISS_INDEX_INFO and index_name not in IMPACT_INDEX_INFO:
        raise ValueError(f'Unrecognized index name {index_name}')
    if index_name in TF_INDEX_INFO:
        target_index = TF_INDEX_INFO[index_name]
    elif index_name in IMPACT_INDEX_INFO:
        target_index = IMPACT_INDEX_INFO[index_name]
    else:
        target_index = FAISS_INDEX_INFO[index_name]
    index_md5 = target_index['md5']
    for url in target_index['urls']:
        local_filename = target_index['filename'] if 'filename' in target_index else None
        try:
            return download_and_unpack_index(url, local_filename=local_filename,
                                             prebuilt=True, md5=index_md5, verbose=verbose)
        except (HTTPError, URLError) as e:
            print(f'Unable to download pre-built index at {url}, trying next URL...')
    raise ValueError(f'Unable to download pre-built index at any known URLs.')


def download_encoded_queries(query_name, force=False, verbose=True, mirror=None):
    if query_name not in QUERY_INFO:
        raise ValueError(f'Unrecognized query name {query_name}')
    query_md5 = QUERY_INFO[query_name]['md5']
    for url in QUERY_INFO[query_name]['urls']:
        try:
            return download_and_unpack_index(url, index_directory='queries', prebuilt=True, md5=query_md5)
        except (HTTPError, URLError) as e:
            print(f'Unable to download encoded query at {url}, trying next URL...')
    raise ValueError(f'Unable to download encoded query at any known URLs.')


def download_encoded_corpus(corpus_name, force=False, verbose=True, mirror=None):
    if corpus_name not in CORPUS_INFO:
        raise ValueError(f'Unrecognized corpus name {corpus_name}')
    corpus_md5 = CORPUS_INFO[corpus_name]['md5']
    for url in CORPUS_INFO[corpus_name]['urls']:
        local_filename = CORPUS_INFO[corpus_name]['filename'] if 'filename' in CORPUS_INFO[corpus_name] else None
        try:
            return download_and_unpack_index(url, local_filename=local_filename, index_directory='corpus', prebuilt=True, md5=corpus_md5)
        except (HTTPError, URLError) as e:
            print(f'Unable to download encoded corpus at {url}, trying next URL...')
    raise ValueError(f'Unable to download encoded corpus at any known URLs.')


def download_evaluation_script(evaluation_name, force=False, verbose=True, mirror=None):
    if evaluation_name not in EVALUATION_INFO:
        raise ValueError(f'Unrecognized evaluation name {evaluation_name}')
    for url in EVALUATION_INFO[evaluation_name]['urls']:
        try:
            save_dir = os.path.join(get_cache_home(), 'eval')
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            return download_url(url, save_dir=save_dir)
        except HTTPError:
            print(f'Unable to download evaluation script at {url}, trying next URL...')
    raise ValueError(f'Unable to download evaluation script at any known URLs.')


def get_sparse_index(index_name):
    if index_name not in FAISS_INDEX_INFO:
        raise ValueError(f'Unrecognized index name {index_name}')
    return FAISS_INDEX_INFO[index_name]["texts"]