Repository: castorini/pyserini

Issue:
Refactoring/Upgrade: Exposing topics and other fixes
This is the package currently deployed at TestPyPI: https://test.pypi.org/project/pyserini/0.6.1.post0/

The fatjar is based on `0ed488` in Anserini, dated `Sun Nov 17 19:08:26 2019 -0500`.

Working colab notebook at: https://colab.research.google.com/drive/1nY1bjwop3Enygrks-208EilhiCYmYKT6

Still need to update docs, but code itself is ready for CR I think.

I'd also appreciate opinions on publishing SNAPSHOT versions of the fatjar that aren't in Maven... Current package is on TestPyPI, so it really doesn't matter... but any strong opinions on publishing in official PyPI?

The tradeoff is between iteration speed vs. "doing it right". PyPI packages are supposed to be permanent, so depending on a snapshot jar is janky.


Discussion:
I think we can hold off on publishing on PyPI for now? We can re-consider later when the Anserini/Pyserini connection is more mature and doesn't require as frequent updates?
Okay, and continue to push out interim releases on TestPyPI for those on the real bleeding edge?
Yes, sounds like a plan!
I'll let @emmileaf take a look before merging...
LGTM!

Resolving patch:
diff --git a/pyserini/search/_base.py b/pyserini/search/_base.py
--- a/pyserini/search/_base.py
+++ b/pyserini/search/_base.py
@@ -208,6 +208,10 @@ def get_topics(collection_name):
         topics = JTopicReader.getTopicsWithStringIds(JTopics.DPR_SQUAD_TEST)
     elif collection_name == 'dpr-curated-test':
         topics = JTopicReader.getTopicsWithStringIds(JTopics.DPR_CURATED_TEST)
+    elif collection_name == 'nq-dev':
+        topics = JTopicReader.getTopicsWithStringIds(JTopics.NQ_DEV)
+    elif collection_name == 'nq-test':
+        topics = JTopicReader.getTopicsWithStringIds(JTopics.NQ_TEST)
     else:
         raise ValueError(f'Topic {collection_name} Not Found')
     t = {}
diff --git a/scripts/dkrr/convert_dkrr_format_to_topic_file.py b/scripts/dkrr/convert_dkrr_format_to_topic_file.py
new file mode 100644
--- /dev/null
+++ b/scripts/dkrr/convert_dkrr_format_to_topic_file.py
@@ -0,0 +1,46 @@
+#
+# 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 argparse
+import json
+
+
+def load_data(data_path):
+    with open(data_path, 'r') as fin:
+        data = json.load(fin)
+    examples = []
+    for k, example in enumerate(data):
+        if not 'id' in example:
+            example['id'] = k
+
+        example['answers'] = list(map(lambda answer: answer.replace('\xa0', ' ').replace('"','""'), example['answers']))
+        examples.append(example)
+    return examples
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', type=str, metavar='path',
+                        help="Path to data json file")
+    parser.add_argument('--output', type=str, metavar='path',
+                        help="Path to topic json file")
+    args = parser.parse_args()
+
+    topic = {}
+    with open(args.output, 'w') as fout:
+        for idx, example in enumerate(load_data(args.input)):
+            fout.write(f"{example['question']}\t{example['answers']}\n")
diff --git a/scripts/dkrr/encode_queries.py b/scripts/dkrr/encode_queries.py
new file mode 100644
--- /dev/null
+++ b/scripts/dkrr/encode_queries.py
@@ -0,0 +1,73 @@
+#
+# 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 argparse
+
+import pandas as pd
+
+from tqdm import tqdm
+from pyserini.query_iterator import get_query_iterator, TopicsFormat
+from transformers import BertModel,  BertTokenizerFast
+import torch
+
+class DkrrDprQueryEncoder():
+
+    def __init__(self, encoder: str = None, device: str = 'cpu', prefix: str = "question:"):
+        self.device = device
+        self.model = BertModel.from_pretrained(encoder)
+        self.model.to(self.device)
+        self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
+        self.prefix = prefix
+
+    @staticmethod
+    def _mean_pooling(model_output, attention_mask):
+        model_output = model_output[0].masked_fill(1 - attention_mask[:, :, None], 0.)
+        model_output = torch.sum(model_output, dim=1) / torch.clamp(torch.sum(attention_mask, dim=1), min=1e-9)[:, None]
+        return model_output.flatten()
+
+    def encode(self, query: str):
+        if self.prefix:
+            query = f'{self.prefix} {query}'
+        inputs = self.tokenizer(query, return_tensors='pt', max_length=40, padding="max_length")
+        inputs.to(self.device)
+        outputs = self.model(input_ids=inputs["input_ids"],
+                            attention_mask=inputs["attention_mask"])
+        embeddings = self._mean_pooling(outputs, inputs['attention_mask']).detach().cpu().numpy()
+        return embeddings.flatten()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--topics', type=str, metavar='topic_name', required=True,
+                        help="Name of topics.")
+    parser.add_argument('--encoder', type=str, help='encoder name or path',
+                        default='facebook/dpr-question_encoder-multiset-base', required=False)
+    parser.add_argument('--output', type=str, help='path to store query embeddings', required=True)
+    parser.add_argument('--device', type=str,
+                        help='device cpu or cuda [cuda:0, cuda:1...]', default='cpu', required=False)
+    args = parser.parse_args()
+
+    query_iterator = get_query_iterator(args.topics, TopicsFormat(TopicsFormat.DEFAULT.value))
+    topics = query_iterator.topics
+    
+    encoder = DkrrDprQueryEncoder(args.encoder, args.device)
+
+    embeddings = {'id': [], 'text': [], 'embedding': []}
+    for index, (topic_id, text) in enumerate(tqdm(query_iterator, total=len(topics.keys()))):
+        embeddings['id'].append(topic_id)
+        embeddings['text'].append(text)
+        embeddings['embedding'].append(encoder.encode(text))
+    embeddings = pd.DataFrame(embeddings)
+    embeddings.to_pickle(args.output)


Test patch: