Repository: castorini/pyserini

Issue:
Dump out tf-idf document vectors in pyserini
Is it possible to dump out tf-idf document vectors for retrieved documents in pyserini?


Discussion:
hi @zeynepakkalyoncu - let's see if we can make this happen?
So, if we just want td-idf values, it'll be fairly easy, we can just compute the td-idf values in Python.

However, if we want accurate BM25 values, we'll have to do something like this:
https://github.com/castorini/anserini/blob/master/src/test/java/io/anserini/integration/IndexerTest.java#L264

... which will be slow.
Could you please describe your solution to compute tf-idf values in python?
hi @alirezamirzaeiyan 

We just need to port this over to Java:
https://github.com/castorini/anserini/blob/master/src/main/java/io/anserini/index/IndexUtils.java#L219

It's on @zeynepakkalyoncu 's todo list, but if you want to do it yourself, code contributions are welcome!

Resolving patch:
diff --git a/pyserini/index/pyutils.py b/pyserini/index/pyutils.py
new file mode 100644
--- /dev/null
+++ b/pyserini/index/pyutils.py
@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+#
+# Pyserini: Python interface to the Anserini IR toolkit built on Lucene
+#
+# 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.
+
+'''
+Module for providing python interface to Anserini index utils
+'''
+from ..pyclass import JIndexUtils, JDocumentVectorWeight, JString
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class IndexUtils:
+    '''
+    Wrapper class for Anserini's IndexUtils.
+
+    Parameters
+    ----------
+    index_dir : str
+        Path to Lucene index directory
+    '''
+
+    def __init__(self, index_dir):
+        self.object = JIndexUtils(JString(index_dir))
+
+    class DocumentVectorWeight:
+        NONE = JDocumentVectorWeight.NONE
+        TF_IDF = JDocumentVectorWeight.TF_IDF
+
+    def dump_document_vectors(self, reqDocidsPath: str, weight: DocumentVectorWeight):
+        '''
+        Parameters
+        ----------
+        reqDocidsPath : str
+            dumps the document vector for all documents in reqDocidsPath
+        weight : DocumentVectorWeight
+            the weight for dumped document vector(s)
+        '''
+        self.object.dumpDocumentVectors(reqDocidsPath, weight)
diff --git a/pyserini/pyclass.py b/pyserini/pyclass.py
--- a/pyserini/pyclass.py
+++ b/pyserini/pyclass.py
@@ -23,7 +23,7 @@
 from .setup import configure_classpath, os
 # If the environment variable isn't defined, look in the current directory.
 configure_classpath(os.environ['ANSERINI_CLASSPATH'] if 'ANSERINI_CLASSPATH' in os.environ else os.path.join(os.path.split(__file__)[0], 'resources/jars/'))
-    
+
 from jnius import autoclass, cast
 from enum import Enum
 
@@ -44,6 +44,10 @@
 JTopics = autoclass('io.anserini.search.topicreader.TopicReader$Topics')
 JTopicReader = autoclass('io.anserini.search.topicreader.TopicReader')
 
+## IndexUtils
+JIndexUtils = autoclass('io.anserini.index.IndexUtils')
+JDocumentVectorWeight = autoclass('io.anserini.index.IndexUtils$DocumentVectorWeight')
+
 ### Generator
 
 class JIndexHelpers:


Test patch: