Mituuu commited on
Commit
3f002dd
·
verified ·
1 Parent(s): 7965a28

Upload 11 files

Browse files
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """# TASK3: a search-engine demo based on Huggingface space (4 points)
2
+
3
+ ## TASK3.1: create the gradio app (2 point)
4
+
5
+ Create a gradio app to demo the BM25 search engine index on SciQ. The app should have a single input variable for the query (of type `str`) and a single output variable for the returned ranking (of type `List[Hit]` in the code below).
6
+ """
7
+ from typing import TypedDict, Optional, List
8
+ import gradio as gr
9
+ from copy_of_hw1 import BM25Retriever
10
+
11
+ class Hit(TypedDict):
12
+ cid: str
13
+ score: float
14
+ text: str
15
+
16
+ demo: Optional[gr.Interface] = None # Assign your gradio demo to this variable
17
+ return_type = List[Hit]
18
+
19
+ ## YOUR_CODE_STARTS_HERE
20
+ def hits(query):
21
+ Hits = []
22
+
23
+ bm25_retriever = BM25Retriever(index_dir="output/bm25_index")
24
+ retrieved = bm25_retriever.retrieve(query)
25
+ for cid in retrieved.keys():
26
+ docid = bm25_retriever.index.cid2docid[cid]
27
+ doc_text = bm25_retriever.index.doc_texts[docid]
28
+ Hits.append(Hit(cid=cid, score=retrieved[cid], text=doc_text))
29
+ return Hits
30
+
31
+ demo = gr.Interface(
32
+ fn=hits,
33
+ inputs=["text"],
34
+ outputs=["text"],
35
+ )
36
+ ## YOUR_CODE_ENDS_HERE
37
+ demo.launch()
copy_of_hw1.py ADDED
@@ -0,0 +1,295 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+ from dataclasses import dataclass
3
+ import pickle
4
+ import os
5
+ from typing import Iterable, Callable, List, Dict, Optional, Type, TypeVar
6
+ from nlp4web_codebase.ir.data_loaders.dm import Document
7
+ from collections import Counter
8
+ import tqdm
9
+ import re
10
+ import nltk
11
+ nltk.download("stopwords", quiet=True)
12
+ from nltk.corpus import stopwords as nltk_stopwords
13
+
14
+ LANGUAGE = "english"
15
+ word_splitter = re.compile(r"(?u)\b\w\w+\b").findall
16
+ stopwords = set(nltk_stopwords.words(LANGUAGE))
17
+
18
+
19
+ def word_splitting(text: str) -> List[str]:
20
+ return word_splitter(text.lower())
21
+
22
+ def lemmatization(words: List[str]) -> List[str]:
23
+ return words # We ignore lemmatization here for simplicity
24
+
25
+ def simple_tokenize(text: str) -> List[str]:
26
+ words = word_splitting(text)
27
+ tokenized = list(filter(lambda w: w not in stopwords, words))
28
+ tokenized = lemmatization(tokenized)
29
+ return tokenized
30
+
31
+ T = TypeVar("T", bound="InvertedIndex")
32
+
33
+ @dataclass
34
+ class PostingList:
35
+ term: str # The term
36
+ docid_postings: List[int] # docid_postings[i] means the docid (int) of the i-th associated posting
37
+ tweight_postings: List[float] # tweight_postings[i] means the term weight (float) of the i-th associated posting
38
+
39
+
40
+ @dataclass
41
+ class InvertedIndex:
42
+ posting_lists: List[PostingList] # docid -> posting_list
43
+ vocab: Dict[str, int]
44
+ cid2docid: Dict[str, int] # collection_id -> docid
45
+ collection_ids: List[str] # docid -> collection_id
46
+ doc_texts: Optional[List[str]] = None # docid -> document text
47
+
48
+ def save(self, output_dir: str) -> None:
49
+ os.makedirs(output_dir, exist_ok=True)
50
+ with open(os.path.join(output_dir, "index.pkl"), "wb") as f:
51
+ pickle.dump(self, f)
52
+
53
+ @classmethod
54
+ def from_saved(cls: Type[T], saved_dir: str) -> T:
55
+ index = cls(
56
+ posting_lists=[], vocab={}, cid2docid={}, collection_ids=[], doc_texts=None
57
+ )
58
+ with open(os.path.join(saved_dir, "index.pkl"), "rb") as f:
59
+ index = pickle.load(f)
60
+ return index
61
+
62
+
63
+ # The output of the counting function:
64
+ @dataclass
65
+ class Counting:
66
+ posting_lists: List[PostingList]
67
+ vocab: Dict[str, int]
68
+ cid2docid: Dict[str, int]
69
+ collection_ids: List[str]
70
+ dfs: List[int] # tid -> df
71
+ dls: List[int] # docid -> doc length
72
+ avgdl: float
73
+ nterms: int
74
+ doc_texts: Optional[List[str]] = None
75
+
76
+ def run_counting(
77
+ documents: Iterable[Document],
78
+ tokenize_fn: Callable[[str], List[str]] = simple_tokenize,
79
+ store_raw: bool = True, # store the document text in doc_texts
80
+ ndocs: Optional[int] = None,
81
+ show_progress_bar: bool = True,
82
+ ) -> Counting:
83
+ """Counting TFs, DFs, doc_lengths, etc."""
84
+ posting_lists: List[PostingList] = []
85
+ vocab: Dict[str, int] = {}
86
+ cid2docid: Dict[str, int] = {}
87
+ collection_ids: List[str] = []
88
+ dfs: List[int] = [] # tid -> df
89
+ dls: List[int] = [] # docid -> doc length
90
+ nterms: int = 0
91
+ doc_texts: Optional[List[str]] = []
92
+ for doc in tqdm.tqdm(
93
+ documents,
94
+ desc="Counting",
95
+ total=ndocs,
96
+ disable=not show_progress_bar,
97
+ ):
98
+ if doc.collection_id in cid2docid:
99
+ continue
100
+ collection_ids.append(doc.collection_id)
101
+ docid = cid2docid.setdefault(doc.collection_id, len(cid2docid))
102
+ toks = tokenize_fn(doc.text)
103
+ tok2tf = Counter(toks)
104
+ dls.append(sum(tok2tf.values()))
105
+ for tok, tf in tok2tf.items():
106
+ nterms += tf
107
+ tid = vocab.get(tok, None)
108
+ if tid is None:
109
+ posting_lists.append(
110
+ PostingList(term=tok, docid_postings=[], tweight_postings=[])
111
+ )
112
+ tid = vocab.setdefault(tok, len(vocab))
113
+ posting_lists[tid].docid_postings.append(docid)
114
+ posting_lists[tid].tweight_postings.append(tf)
115
+ if tid < len(dfs):
116
+ dfs[tid] += 1
117
+ else:
118
+ dfs.append(0)
119
+ if store_raw:
120
+ doc_texts.append(doc.text)
121
+ else:
122
+ doc_texts = None
123
+ return Counting(
124
+ posting_lists=posting_lists,
125
+ vocab=vocab,
126
+ cid2docid=cid2docid,
127
+ collection_ids=collection_ids,
128
+ dfs=dfs,
129
+ dls=dls,
130
+ avgdl=sum(dls) / len(dls),
131
+ nterms=nterms,
132
+ doc_texts=doc_texts,
133
+ )
134
+
135
+ from nlp4web_codebase.ir.data_loaders.sciq import load_sciq
136
+ sciq = load_sciq()
137
+ counting = run_counting(documents=iter(sciq.corpus), ndocs=len(sciq.corpus))
138
+
139
+ from dataclasses import asdict, dataclass
140
+ import math
141
+ import os
142
+ from typing import Iterable, List, Optional, Type
143
+ import tqdm
144
+ from nlp4web_codebase.ir.data_loaders.dm import Document
145
+
146
+
147
+ @dataclass
148
+ class BM25Index(InvertedIndex):
149
+
150
+ @staticmethod
151
+ def tokenize(text: str) -> List[str]:
152
+ return simple_tokenize(text)
153
+
154
+ @staticmethod
155
+ def cache_term_weights(
156
+ posting_lists: List[PostingList],
157
+ total_docs: int,
158
+ avgdl: float,
159
+ dfs: List[int],
160
+ dls: List[int],
161
+ k1: float,
162
+ b: float,
163
+ ) -> None:
164
+ """Compute term weights and caching"""
165
+
166
+ N = total_docs
167
+ for tid, posting_list in enumerate(
168
+ tqdm.tqdm(posting_lists, desc="Regularizing TFs")
169
+ ):
170
+ idf = BM25Index.calc_idf(df=dfs[tid], N=N)
171
+ for i in range(len(posting_list.docid_postings)):
172
+ docid = posting_list.docid_postings[i]
173
+ tf = posting_list.tweight_postings[i]
174
+ dl = dls[docid]
175
+ regularized_tf = BM25Index.calc_regularized_tf(
176
+ tf=tf, dl=dl, avgdl=avgdl, k1=k1, b=b
177
+ )
178
+ posting_list.tweight_postings[i] = regularized_tf * idf
179
+
180
+ @staticmethod
181
+ def calc_regularized_tf(
182
+ tf: int, dl: float, avgdl: float, k1: float, b: float
183
+ ) -> float:
184
+ return tf / (tf + k1 * (1 - b + b * dl / avgdl))
185
+
186
+ @staticmethod
187
+ def calc_idf(df: int, N: int):
188
+ return math.log(1 + (N - df + 0.5) / (df + 0.5))
189
+
190
+ @classmethod
191
+ def build_from_documents(
192
+ cls: Type[BM25Index],
193
+ documents: Iterable[Document],
194
+ store_raw: bool = True,
195
+ output_dir: Optional[str] = None,
196
+ ndocs: Optional[int] = None,
197
+ show_progress_bar: bool = True,
198
+ k1: float = 0.9,
199
+ b: float = 0.4,
200
+ ) -> BM25Index:
201
+ # Counting TFs, DFs, doc_lengths, etc.:
202
+ counting = run_counting(
203
+ documents=documents,
204
+ tokenize_fn=BM25Index.tokenize,
205
+ store_raw=store_raw,
206
+ ndocs=ndocs,
207
+ show_progress_bar=show_progress_bar,
208
+ )
209
+
210
+ # Compute term weights and caching:
211
+ posting_lists = counting.posting_lists
212
+ total_docs = len(counting.cid2docid)
213
+ BM25Index.cache_term_weights(
214
+ posting_lists=posting_lists,
215
+ total_docs=total_docs,
216
+ avgdl=counting.avgdl,
217
+ dfs=counting.dfs,
218
+ dls=counting.dls,
219
+ k1=k1,
220
+ b=b,
221
+ )
222
+
223
+ # Assembly and save:
224
+ index = BM25Index(
225
+ posting_lists=posting_lists,
226
+ vocab=counting.vocab,
227
+ cid2docid=counting.cid2docid,
228
+ collection_ids=counting.collection_ids,
229
+ doc_texts=counting.doc_texts,
230
+ )
231
+ return index
232
+
233
+
234
+ from nlp4web_codebase.ir.models import BaseRetriever
235
+ from typing import Type
236
+ from abc import abstractmethod
237
+
238
+
239
+ class BaseInvertedIndexRetriever(BaseRetriever):
240
+
241
+ @property
242
+ @abstractmethod
243
+ def index_class(self) -> Type[InvertedIndex]:
244
+ pass
245
+
246
+ def __init__(self, index_dir: str) -> None:
247
+ self.index = self.index_class.from_saved(index_dir)
248
+
249
+ def get_term_weights(self, query: str, cid: str) -> Dict[str, float]:
250
+ toks = self.index.tokenize(query)
251
+ target_docid = self.index.cid2docid[cid]
252
+ term_weights = {}
253
+ for tok in toks:
254
+ if tok not in self.index.vocab:
255
+ continue
256
+ tid = self.index.vocab[tok]
257
+ posting_list = self.index.posting_lists[tid]
258
+ for docid, tweight in zip(
259
+ posting_list.docid_postings, posting_list.tweight_postings
260
+ ):
261
+ if docid == target_docid:
262
+ term_weights[tok] = tweight
263
+ break
264
+ return term_weights
265
+
266
+ def score(self, query: str, cid: str) -> float:
267
+ return sum(self.get_term_weights(query=query, cid=cid).values())
268
+
269
+ def retrieve(self, query: str, topk: int = 10) -> Dict[str, float]:
270
+ toks = self.index.tokenize(query)
271
+ docid2score: Dict[int, float] = {}
272
+ for tok in toks:
273
+ if tok not in self.index.vocab:
274
+ continue
275
+ tid = self.index.vocab[tok]
276
+ posting_list = self.index.posting_lists[tid]
277
+ for docid, tweight in zip(
278
+ posting_list.docid_postings, posting_list.tweight_postings
279
+ ):
280
+ docid2score.setdefault(docid, 0)
281
+ docid2score[docid] += tweight
282
+ docid2score = dict(
283
+ sorted(docid2score.items(), key=lambda pair: pair[1], reverse=True)[:topk]
284
+ )
285
+ return {
286
+ self.index.collection_ids[docid]: score
287
+ for docid, score in docid2score.items()
288
+ }
289
+
290
+
291
+ class BM25Retriever(BaseInvertedIndexRetriever):
292
+
293
+ @property
294
+ def index_class(self) -> Type[BM25Index]:
295
+ return BM25Index
nlp4web_codebase/__init__.py ADDED
File without changes
nlp4web_codebase/ir/__init__.py ADDED
File without changes
nlp4web_codebase/ir/analysis.py ADDED
@@ -0,0 +1,160 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from typing import Dict, List, Optional, Protocol
3
+ import pandas as pd
4
+ import tqdm
5
+ import ujson
6
+ from nlp4web_codebase.ir.data_loaders import IRDataset
7
+
8
+
9
+ def round_dict(obj: Dict[str, float], ndigits: int = 4) -> Dict[str, float]:
10
+ return {k: round(v, ndigits=ndigits) for k, v in obj.items()}
11
+
12
+
13
+ def sort_dict(obj: Dict[str, float], reverse: bool = True) -> Dict[str, float]:
14
+ return dict(sorted(obj.items(), key=lambda pair: pair[1], reverse=reverse))
15
+
16
+
17
+ def save_ranking_results(
18
+ output_dir: str,
19
+ query_ids: List[str],
20
+ rankings: List[Dict[str, float]],
21
+ query_performances_lists: List[Dict[str, float]],
22
+ cid2tweights_lists: Optional[List[Dict[str, Dict[str, float]]]] = None,
23
+ ):
24
+ os.makedirs(output_dir, exist_ok=True)
25
+ output_path = os.path.join(output_dir, "ranking_results.jsonl")
26
+ rows = []
27
+ for i, (query_id, ranking, query_performances) in enumerate(
28
+ zip(query_ids, rankings, query_performances_lists)
29
+ ):
30
+ row = {
31
+ "query_id": query_id,
32
+ "ranking": round_dict(ranking),
33
+ "query_performances": round_dict(query_performances),
34
+ "cid2tweights": {},
35
+ }
36
+ if cid2tweights_lists is not None:
37
+ row["cid2tweights"] = {
38
+ cid: round_dict(tws) for cid, tws in cid2tweights_lists[i].items()
39
+ }
40
+ rows.append(row)
41
+ pd.DataFrame(rows).to_json(
42
+ output_path,
43
+ orient="records",
44
+ lines=True,
45
+ )
46
+
47
+
48
+ class TermWeightingFunction(Protocol):
49
+ def __call__(self, query: str, cid: str) -> Dict[str, float]: ...
50
+
51
+
52
+ def compare(
53
+ dataset: IRDataset,
54
+ results_path1: str,
55
+ results_path2: str,
56
+ output_dir: str,
57
+ main_metric: str = "recip_rank",
58
+ system1: Optional[str] = None,
59
+ system2: Optional[str] = None,
60
+ term_weighting_fn1: Optional[TermWeightingFunction] = None,
61
+ term_weighting_fn2: Optional[TermWeightingFunction] = None,
62
+ ) -> None:
63
+ os.makedirs(output_dir, exist_ok=True)
64
+ df1 = pd.read_json(results_path1, orient="records", lines=True)
65
+ df2 = pd.read_json(results_path2, orient="records", lines=True)
66
+ assert len(df1) == len(df2)
67
+ all_qrels = {}
68
+ for split in dataset.split2qrels:
69
+ all_qrels.update(dataset.get_qrels_dict(split))
70
+ qid2query = {query.query_id: query for query in dataset.queries}
71
+ cid2doc = {doc.collection_id: doc for doc in dataset.corpus}
72
+ diff_col = f"{main_metric}:qp1-qp2"
73
+ merged = pd.merge(df1, df2, on="query_id", how="outer")
74
+ rows = []
75
+ for _, example in tqdm.tqdm(merged.iterrows(), desc="Comparing", total=len(merged)):
76
+ docs = {cid: cid2doc[cid].text for cid in dict(example["ranking_x"])}
77
+ docs.update({cid: cid2doc[cid].text for cid in dict(example["ranking_y"])})
78
+ query_id = example["query_id"]
79
+ row = {
80
+ "query_id": query_id,
81
+ "query": qid2query[query_id].text,
82
+ diff_col: example["query_performances_x"][main_metric]
83
+ - example["query_performances_y"][main_metric],
84
+ "ranking1": ujson.dumps(example["ranking_x"], indent=4),
85
+ "ranking2": ujson.dumps(example["ranking_y"], indent=4),
86
+ "docs": ujson.dumps(docs, indent=4),
87
+ "query_performances1": ujson.dumps(
88
+ example["query_performances_x"], indent=4
89
+ ),
90
+ "query_performances2": ujson.dumps(
91
+ example["query_performances_y"], indent=4
92
+ ),
93
+ "qrels": ujson.dumps(all_qrels[query_id], indent=4),
94
+ }
95
+ if term_weighting_fn1 is not None and term_weighting_fn2 is not None:
96
+ all_cids = set(example["ranking_x"]) | set(example["ranking_y"])
97
+ cid2tweights1 = {}
98
+ cid2tweights2 = {}
99
+ ranking1 = {}
100
+ ranking2 = {}
101
+ for cid in all_cids:
102
+ tweights1 = term_weighting_fn1(query=qid2query[query_id].text, cid=cid)
103
+ tweights2 = term_weighting_fn2(query=qid2query[query_id].text, cid=cid)
104
+ ranking1[cid] = sum(tweights1.values())
105
+ ranking2[cid] = sum(tweights2.values())
106
+ cid2tweights1[cid] = tweights1
107
+ cid2tweights2[cid] = tweights2
108
+ ranking1 = sort_dict(ranking1)
109
+ ranking2 = sort_dict(ranking2)
110
+ row["ranking1"] = ujson.dumps(ranking1, indent=4)
111
+ row["ranking2"] = ujson.dumps(ranking2, indent=4)
112
+ cid2tweights1 = {cid: cid2tweights1[cid] for cid in ranking1}
113
+ cid2tweights2 = {cid: cid2tweights2[cid] for cid in ranking2}
114
+ row["cid2tweights1"] = ujson.dumps(cid2tweights1, indent=4)
115
+ row["cid2tweights2"] = ujson.dumps(cid2tweights2, indent=4)
116
+ rows.append(row)
117
+ table = pd.DataFrame(rows).sort_values(by=diff_col, ascending=False)
118
+ output_path = os.path.join(output_dir, f"compare-{system1}_vs_{system2}.tsv")
119
+ table.to_csv(output_path, sep="\t", index=False)
120
+
121
+
122
+ # if __name__ == "__main__":
123
+ # # python -m lecture2.bm25.analysis
124
+ # from nlp4web_codebase.ir.data_loaders.sciq import load_sciq
125
+ # from lecture2.bm25.bm25_retriever import BM25Retriever
126
+ # from lecture2.bm25.tfidf_retriever import TFIDFRetriever
127
+ # import numpy as np
128
+
129
+ # sciq = load_sciq()
130
+ # system1 = "bm25"
131
+ # system2 = "tfidf"
132
+ # results_path1 = f"output/sciq-{system1}/results/ranking_results.jsonl"
133
+ # results_path2 = f"output/sciq-{system2}/results/ranking_results.jsonl"
134
+ # index_dir1 = f"output/sciq-{system1}"
135
+ # index_dir2 = f"output/sciq-{system2}"
136
+ # compare(
137
+ # dataset=sciq,
138
+ # results_path1=results_path1,
139
+ # results_path2=results_path2,
140
+ # output_dir=f"output/sciq-{system1}_vs_{system2}",
141
+ # system1=system1,
142
+ # system2=system2,
143
+ # term_weighting_fn1=BM25Retriever(index_dir1).get_term_weights,
144
+ # term_weighting_fn2=TFIDFRetriever(index_dir2).get_term_weights,
145
+ # )
146
+
147
+ # # bias on #shared_terms of TFIDF:
148
+ # df1 = pd.read_json(results_path1, orient="records", lines=True)
149
+ # df2 = pd.read_json(results_path2, orient="records", lines=True)
150
+ # merged = pd.merge(df1, df2, on="query_id", how="outer")
151
+ # nterms1 = []
152
+ # nterms2 = []
153
+ # for _, row in merged.iterrows():
154
+ # nterms1.append(len(list(dict(row["cid2tweights_x"]).values())[0]))
155
+ # nterms2.append(len(list(dict(row["cid2tweights_y"]).values())[0]))
156
+ # percentiles = (5, 25, 50, 75, 95)
157
+ # print(system1, np.percentile(nterms1, percentiles), np.mean(nterms1).round(2))
158
+ # print(system2, np.percentile(nterms2, percentiles), np.mean(nterms2).round(2))
159
+ # # bm25 [ 3. 4. 5. 7. 11.] 5.64
160
+ # # tfidf [1. 2. 3. 5. 9.] 3.58
nlp4web_codebase/ir/data_loaders/__init__.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass
2
+ from enum import Enum
3
+ from typing import Dict, List
4
+ from nlp4web_codebase.ir.data_loaders.dm import Document, Query, QRel
5
+
6
+
7
+ class Split(str, Enum):
8
+ train = "train"
9
+ dev = "dev"
10
+ test = "test"
11
+
12
+
13
+ @dataclass
14
+ class IRDataset:
15
+ corpus: List[Document]
16
+ queries: List[Query]
17
+ split2qrels: Dict[Split, List[QRel]]
18
+
19
+ def get_stats(self) -> Dict[str, int]:
20
+ stats = {"|corpus|": len(self.corpus), "|queries|": len(self.queries)}
21
+ for split, qrels in self.split2qrels.items():
22
+ stats[f"|qrels-{split}|"] = len(qrels)
23
+ return stats
24
+
25
+ def get_qrels_dict(self, split: Split) -> Dict[str, Dict[str, int]]:
26
+ qrels_dict = {}
27
+ for qrel in self.split2qrels[split]:
28
+ qrels_dict.setdefault(qrel.query_id, {})
29
+ qrels_dict[qrel.query_id][qrel.collection_id] = qrel.relevance
30
+ return qrels_dict
31
+
32
+ def get_split_queries(self, split: Split) -> List[Query]:
33
+ qrels = self.split2qrels[split]
34
+ qids = {qrel.query_id for qrel in qrels}
35
+ return list(filter(lambda query: query.query_id in qids, self.queries))
nlp4web_codebase/ir/data_loaders/dm.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass
2
+ from typing import Optional
3
+
4
+
5
+ @dataclass
6
+ class Document:
7
+ collection_id: str
8
+ text: str
9
+
10
+
11
+ @dataclass
12
+ class Query:
13
+ query_id: str
14
+ text: str
15
+
16
+
17
+ @dataclass
18
+ class QRel:
19
+ query_id: str
20
+ collection_id: str
21
+ relevance: int
22
+ answer: Optional[str] = None
nlp4web_codebase/ir/data_loaders/sciq.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List
2
+ from nlp4web_codebase.ir.data_loaders import IRDataset, Split
3
+ from nlp4web_codebase.ir.data_loaders.dm import Document, Query, QRel
4
+ from datasets import load_dataset
5
+ import joblib
6
+
7
+
8
+ @(joblib.Memory(".cache").cache)
9
+ def load_sciq(verbose: bool = False) -> IRDataset:
10
+ train = load_dataset("allenai/sciq", split="train")
11
+ validation = load_dataset("allenai/sciq", split="validation")
12
+ test = load_dataset("allenai/sciq", split="test")
13
+ data = {Split.train: train, Split.dev: validation, Split.test: test}
14
+
15
+ # Each duplicated record is the same to each other:
16
+ df = train.to_pandas() + validation.to_pandas() + test.to_pandas()
17
+ for question, group in df.groupby("question"):
18
+ assert len(set(group["support"].tolist())) == len(group)
19
+ assert len(set(group["correct_answer"].tolist())) == len(group)
20
+
21
+ # Build:
22
+ corpus = []
23
+ queries = []
24
+ split2qrels: Dict[str, List[dict]] = {}
25
+ question2id = {}
26
+ support2id = {}
27
+ for split, rows in data.items():
28
+ if verbose:
29
+ print(f"|raw_{split}|", len(rows))
30
+ split2qrels[split] = []
31
+ for i, row in enumerate(rows):
32
+ example_id = f"{split}-{i}"
33
+ support: str = row["support"]
34
+ if len(support.strip()) == 0:
35
+ continue
36
+ question = row["question"]
37
+ if len(support.strip()) == 0:
38
+ continue
39
+ if support in support2id:
40
+ continue
41
+ else:
42
+ support2id[support] = example_id
43
+ if question in question2id:
44
+ continue
45
+ else:
46
+ question2id[question] = example_id
47
+ doc = {"collection_id": example_id, "text": support}
48
+ query = {"query_id": example_id, "text": row["question"]}
49
+ qrel = {
50
+ "query_id": example_id,
51
+ "collection_id": example_id,
52
+ "relevance": 1,
53
+ "answer": row["correct_answer"],
54
+ }
55
+ corpus.append(Document(**doc))
56
+ queries.append(Query(**query))
57
+ split2qrels[split].append(QRel(**qrel))
58
+
59
+ # Assembly and return:
60
+ return IRDataset(corpus=corpus, queries=queries, split2qrels=split2qrels)
61
+
62
+
63
+ if __name__ == "__main__":
64
+ # python -m nlp4web_codebase.ir.data_loaders.sciq
65
+ import ujson
66
+ import time
67
+
68
+ start = time.time()
69
+ dataset = load_sciq(verbose=True)
70
+ print(f"Loading costs: {time.time() - start}s")
71
+ print(ujson.dumps(dataset.get_stats(), indent=4))
72
+ # ________________________________________________________________________________
73
+ # [Memory] Calling __main__--home-kwang-research-nlp4web-ir-exercise-nlp4web-nlp4web-ir-data_loaders-sciq.load_sciq...
74
+ # load_sciq(verbose=True)
75
+ # |raw_train| 11679
76
+ # |raw_dev| 1000
77
+ # |raw_test| 1000
78
+ # ________________________________________________________load_sciq - 7.3s, 0.1min
79
+ # Loading costs: 7.260092735290527s
80
+ # {
81
+ # "|corpus|": 12160,
82
+ # "|queries|": 12160,
83
+ # "|qrels-train|": 10409,
84
+ # "|qrels-dev|": 875,
85
+ # "|qrels-test|": 876
86
+ # }
nlp4web_codebase/ir/models/__init__.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from abc import ABC, abstractmethod
2
+ from typing import Any, Dict, Type
3
+
4
+
5
+ class BaseRetriever(ABC):
6
+
7
+ @property
8
+ @abstractmethod
9
+ def index_class(self) -> Type[Any]:
10
+ pass
11
+
12
+ def get_term_weights(self, query: str, cid: str) -> Dict[str, float]:
13
+ raise NotImplementedError
14
+
15
+ @abstractmethod
16
+ def score(self, query: str, cid: str) -> float:
17
+ pass
18
+
19
+ @abstractmethod
20
+ def retrieve(self, query: str, topk: int = 10) -> Dict[str, float]:
21
+ pass
output/bm25_index/index.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8258988207556a8feb038babf941dccdd2aef9b3cf78eb44bef8a6341c5029f7
3
+ size 11624459
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ joblib
2
+ nltk