pyserini/query_iterator.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 os
import json
from abc import ABC, abstractmethod
from enum import Enum, unique
from pathlib import Path

from pyserini.search import get_topics, get_topics_with_reader
from pyserini.util import download_url, get_cache_home
from pyserini.external_query_info import KILT_QUERY_INFO
from urllib.error import HTTPError, URLError


@unique
class TopicsFormat(Enum):
    DEFAULT = 'default'
    KILT = 'kilt'


class QueryIterator(ABC):

    PREDEFINED_ORDER = {'msmarco-doc-dev',
                        'msmarco-doc-test',
                        'msmarco-passage-dev-subset',
                        'msmarco-passage-test-subset'}

    def __init__(self, topics: dict, order: list = None):
        self.order = order if order else sorted(topics.keys())
        self.topics = topics

    @abstractmethod
    def get_query(self, id_):
        raise NotImplementedError()

    @classmethod
    @abstractmethod
    def from_topics(cls, topics_path: str):
        raise NotImplementedError()

    def __iter__(self):
        for id_ in self.order:
            yield id_, self.get_query(id_)

    def __len__(self):
        return len(self.topics.keys())

    @staticmethod
    def get_predefined_order(topics_path: str):
        order = None
        normalized_path = Path(topics_path).stem  # get filename w/o extension
        normalized_path = normalized_path.replace('_', '-')

        if normalized_path in QueryIterator.PREDEFINED_ORDER:
            print(f'Using pre-defined topic order for {normalized_path}')
            # Lazy import:
            from pyserini.query_iterator_order_info import QUERY_IDS
            order = QUERY_IDS[topics_path]
        return order


class DefaultQueryIterator(QueryIterator):

    def get_query(self, id_):
        return self.topics[id_].get('title')

    @classmethod
    def from_topics(cls, topics_path: str):
        if os.path.exists(topics_path):
            if topics_path.endswith('.json'):
                with open(topics_path, 'r') as f:
                    topics = json.load(f)
            elif 'beir' in topics_path:
                topics = get_topics_with_reader('io.anserini.search.topicreader.TsvStringTopicReader', topics_path)
            elif topics_path.endswith('.tsv') or topics_path.endswith('.tsv.gz'):
                try:
                    topics = get_topics_with_reader('io.anserini.search.topicreader.TsvIntTopicReader', topics_path)
                except ValueError as e:
                    topics = get_topics_with_reader('io.anserini.search.topicreader.TsvStringTopicReader', topics_path)
            elif topics_path.endswith('.trec'):
                topics = get_topics_with_reader('io.anserini.search.topicreader.TrecTopicReader', topics_path)
            elif 'cacm' in topics_path:
                topics = get_topics_with_reader('io.anserini.search.topicreader.CacmTopicReader', topics_path)
            elif topics_path.endswith('.jsonl'):
                topics = get_topics_with_reader('io.anserini.search.topicreader.JsonStringTopicReader', topics_path)
            else:
                raise NotImplementedError(f"Not sure how to parse {topics_path}. Please specify the file extension.")
        else:
            topics = get_topics(topics_path)
        if not topics:
            raise FileNotFoundError(f'Topic {topics_path} Not Found')
        order = QueryIterator.get_predefined_order(topics_path)
        return cls(topics, order)


class KiltQueryIterator(QueryIterator):

    ENT_START_TOKEN = "[START_ENT]"
    ENT_END_TOKEN = "[END_ENT]"

    def get_query(self, id_):
        datapoint = self.topics[id_]
        query = (
            datapoint["input"]
            .replace(KiltQueryIterator.ENT_START_TOKEN, "")
            .replace(KiltQueryIterator.ENT_END_TOKEN, "")
            .strip()
        )
        return query

    @classmethod
    def from_topics(cls, topics_path: str):
        topics = {}
        order = []
        if not os.path.exists(topics_path):
            # Download if necessary:
            topics_path = cls.download_kilt_topics(topics_path)
        with open(topics_path, 'r') as f:
            for line in f:
                datapoint = json.loads(line)
                topics[datapoint["id"]] = datapoint
                order.append(datapoint["id"])
        return cls(topics, order)

    @classmethod
    def download_kilt_topics(cls, task: str, force=False):
        if task not in KILT_QUERY_INFO:
            raise ValueError(f'Unrecognized query name {task}')
        task = KILT_QUERY_INFO[task]
        md5 = task['md5']
        save_dir = os.path.join(get_cache_home(), 'queries')
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        for url in task['urls']:
            try:
                return download_url(url, save_dir, force=force, md5=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 get_query_iterator(topics_path: str, topics_format: TopicsFormat):
    mapping = {
        TopicsFormat.DEFAULT: DefaultQueryIterator,
        TopicsFormat.KILT: KiltQueryIterator,
    }
    return mapping[topics_format].from_topics(topics_path)