Move data/ contents

- Move InvertedIndex.py to util/
- Move Dataset class to DatasetParser.py

retrieval/DatasetParser.py

@@ -1,11 +1,30 @@
 import os
 
-from retrieval.data.Dataset import Dataset
-from retrieval.data.InvertedIndex import InvertedIndex
+from retrieval.util.InvertedIndex import InvertedIndex
 from retrieval.models.BM25 import BM25
 from retrieval.models.QueryLikelihood import QueryLikelihood
 from retrieval.models.VectorSpace import VectorSpace
-from retrieval.util.FileManager import read_pickle, write_pickle
+from retrieval.util.FileManager import read_pickle, write_pickle, read_tsv
+
+
+class Dataset:
+    def __init__(self, file_name: str):
+        self._rows = read_tsv(file_name)
+
+    def id_mapping(self) -> dict[int, list[int]]:
+        mapping = {}
+        for row in self._rows:
+            qid, pid = int(row[0]), int(row[1])
+            if qid not in mapping:
+                mapping[qid] = []
+            mapping[qid] += [pid]
+        return mapping
+
+    def queries(self) -> dict[int, str]:
+        return {int(row[0]): row[2] for row in self._rows}
+
+    def passages(self) -> dict[int, str]:
+        return {int(row[1]): row[3] for row in self._rows}
 
 
 class DatasetParser:

retrieval/models/BM25.py

@@ -1,6 +1,6 @@
 from math import log
 
-from retrieval.data.InvertedIndex import InvertedIndex
+from retrieval.util.InvertedIndex import InvertedIndex
 from retrieval.models.Model import Model
 
 

retrieval/models/Model.py

@@ -1,4 +1,4 @@
-from retrieval.data.InvertedIndex import InvertedIndex
+from retrieval.util.InvertedIndex import InvertedIndex
 from retrieval.util.TextProcessor import clean
 
 

retrieval/models/QueryLikelihood.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from retrieval.data.InvertedIndex import InvertedIndex
+from retrieval.util.InvertedIndex import InvertedIndex
 from retrieval.models.Model import Model
 
 

retrieval/models/VectorSpace.py

@@ -4,7 +4,7 @@
 import numpy as np
 import numpy.linalg as npl
 
-from retrieval.data.InvertedIndex import InvertedIndex
+from retrieval.util.InvertedIndex import InvertedIndex
 from retrieval.models.Model import Model