t

app.py

@@ -0,0 +1,183 @@
+import gradio as gr
+from typing import TypedDict
+from typing import Dict, List
+from datasets import load_dataset
+import joblib
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class Document:
+    collection_id: str
+    text: str
+
+
+@dataclass
+class Query:
+    query_id: str
+    text: str
+
+
+@dataclass
+class QRel:
+    query_id: str
+    collection_id: str
+    relevance: int
+    answer: Optional[str] = None
+
+
+
+class Split(str, Enum):
+    train = "train"
+    dev = "dev"
+    test = "test"
+
+
+@dataclass
+class IRDataset:
+    corpus: List[Document]
+    queries: List[Query]
+    split2qrels: Dict[Split, List[QRel]]
+
+    def get_stats(self) -> Dict[str, int]:
+        stats = {"|corpus|": len(self.corpus), "|queries|": len(self.queries)}
+        for split, qrels in self.split2qrels.items():
+            stats[f"|qrels-{split}|"] = len(qrels)
+        return stats
+
+    def get_qrels_dict(self, split: Split) -> Dict[str, Dict[str, int]]:
+        qrels_dict = {}
+        for qrel in self.split2qrels[split]:
+            qrels_dict.setdefault(qrel.query_id, {})
+            qrels_dict[qrel.query_id][qrel.collection_id] = qrel.relevance
+        return qrels_dict
+
+    def get_split_queries(self, split: Split) -> List[Query]:
+        qrels = self.split2qrels[split]
+        qids = {qrel.query_id for qrel in qrels}
+        return list(filter(lambda query: query.query_id in qids, self.queries))
+
+
+@(joblib.Memory(".cache").cache)
+def load_sciq(verbose: bool = False) -> IRDataset:
+    train = load_dataset("allenai/sciq", split="train")
+    validation = load_dataset("allenai/sciq", split="validation")
+    test = load_dataset("allenai/sciq", split="test")
+    data = {Split.train: train, Split.dev: validation, Split.test: test}
+
+    # Each duplicated record is the same to each other:
+    df = train.to_pandas() + validation.to_pandas() + test.to_pandas()
+    for question, group in df.groupby("question"):
+        assert len(set(group["support"].tolist())) == len(group)
+        assert len(set(group["correct_answer"].tolist())) == len(group)
+
+    # Build:
+    corpus = []
+    queries = []
+    split2qrels: Dict[str, List[dict]] = {}
+    question2id = {}
+    support2id = {}
+    for split, rows in data.items():
+        if verbose:
+            print(f"|raw_{split}|", len(rows))
+        split2qrels[split] = []
+        for i, row in enumerate(rows):
+            example_id = f"{split}-{i}"
+            support: str = row["support"]
+            if len(support.strip()) == 0:
+                continue
+            question = row["question"]
+            if len(support.strip()) == 0:
+                continue
+            if support in support2id:
+                continue
+            else:
+                support2id[support] = example_id
+            if question in question2id:
+                continue
+            else:
+                question2id[question] = example_id
+            doc = {"collection_id": example_id, "text": support}
+            query = {"query_id": example_id, "text": row["question"]}
+            qrel = {
+                "query_id": example_id,
+                "collection_id": example_id,
+                "relevance": 1,
+                "answer": row["correct_answer"],
+            }
+            corpus.append(Document(**doc))
+            queries.append(Query(**query))
+            split2qrels[split].append(QRel(**qrel))
+
+    # Assembly and return:
+    return IRDataset(corpus=corpus, queries=queries, split2qrels=split2qrels)
+
+
+if __name__ == "__main__":
+    # python -m nlp4web_codebase.ir.data_loaders.sciq
+    import ujson
+    import time
+
+    start = time.time()
+    dataset = load_sciq(verbose=True)
+    print(f"Loading costs: {time.time() - start}s")
+    print(ujson.dumps(dataset.get_stats(), indent=4))
+    # ________________________________________________________________________________
+    # [Memory] Calling __main__--home-kwang-research-nlp4web-ir-exercise-nlp4web-nlp4web-ir-data_loaders-sciq.load_sciq...
+    # load_sciq(verbose=True)
+    # |raw_train| 11679
+    # |raw_dev| 1000
+    # |raw_test| 1000
+    # ________________________________________________________load_sciq - 7.3s, 0.1min
+    # Loading costs: 7.260092735290527s
+    # {
+    #     "|corpus|": 12160,
+    #     "|queries|": 12160,
+    #     "|qrels-train|": 10409,
+    #     "|qrels-dev|": 875,
+    #     "|qrels-test|": 876
+    # }
+
+
+class Hit(TypedDict):
+  cid: str
+  score: float
+  text: str
+
+## YOUR_CODE_STARTS_HERE
+def search(query: str) -> List[Hit]:
+
+    sciq = load_sciq()
+    counting = run_counting(documents=iter(sciq.corpus), ndocs=len(sciq.corpus))
+
+    bm25_index = BM25Index.build_from_documents(
+        documents=iter(sciq.corpus),
+        ndocs=12160,
+        show_progress_bar=True
+    )
+    bm25_index.save("output/bm25_index")
+
+    bm25_retriever = BM25Retriever(index_dir="output/bm25_index")
+
+    results = bm25_retriever.retrieve(query=query)
+
+    hits: List[Hit] = []
+    for cid, score in results.items():
+        docid = bm25_retriever.index.cid2docid[cid]
+        text = bm25_retriever.index.doc_texts[docid]
+        hits.append({"cid": cid, "score": score, "text": text})
+
+    return hits
+## YOUR_CODE_ENDS_HERE
+
+demo: Optional[gr.Interface] = gr.Interface(
+    fn=search,
+    inputs=gr.Textbox(label="Query"),
+    outputs=gr.JSON(label="Results")
+)  # Assign your gradio demo to this variable
+return_type = List[Hit]
+demo.launch()