File size: 4,961 Bytes
cd6775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | from typing import Dict, Tuple
from tqdm.autonotebook import tqdm
import json
import os
import logging
import csv
logger = logging.getLogger(__name__)
class GenericDataLoader:
def __init__(self, data_folder: str = None, prefix: str = None, corpus_file: str = "corpus.jsonl", query_file: str = "queries.jsonl",
qrels_folder: str = "qrels", qrels_file: str = ""):
self.corpus = {}
self.queries = {}
self.qrels = {}
if prefix:
query_file = prefix + "-" + query_file
qrels_folder = prefix + "-" + qrels_folder
self.corpus_file = os.path.join(data_folder, corpus_file) if data_folder else corpus_file
self.query_file = os.path.join(data_folder, query_file) if data_folder else query_file
self.qrels_folder = os.path.join(data_folder, qrels_folder) if data_folder else None
self.qrels_file = qrels_file
@staticmethod
def check(fIn: str, ext: str):
if not os.path.exists(fIn):
raise ValueError("File {} not present! Please provide accurate file.".format(fIn))
if not fIn.endswith(ext):
raise ValueError("File {} must be present with extension {}".format(fIn, ext))
def load_custom(self) -> Tuple[Dict[str, Dict[str, str]], Dict[str, str], Dict[str, Dict[str, int]]]:
self.check(fIn=self.corpus_file, ext="jsonl")
self.check(fIn=self.query_file, ext="jsonl")
self.check(fIn=self.qrels_file, ext="tsv")
if not len(self.corpus):
logger.info("Loading Corpus...")
self._load_corpus()
logger.info("Loaded %d Documents.", len(self.corpus))
logger.info("Doc Example: %s", list(self.corpus.values())[0])
if not len(self.queries):
logger.info("Loading Queries...")
self._load_queries()
if os.path.exists(self.qrels_file):
self._load_qrels()
self.queries = {qid: self.queries[qid] for qid in self.qrels}
logger.info("Loaded %d Queries.", len(self.queries))
logger.info("Query Example: %s", list(self.queries.values())[0])
return self.corpus, self.queries, self.qrels
def load(self, split="test") -> Tuple[Dict[str, Dict[str, str]], Dict[str, str], Dict[str, Dict[str, int]]]:
self.qrels_file = os.path.join(self.qrels_folder, split + ".tsv")
self.check(fIn=self.corpus_file, ext="jsonl")
self.check(fIn=self.query_file, ext="jsonl")
self.check(fIn=self.qrels_file, ext="tsv")
if not len(self.corpus):
logger.info("Loading Corpus...")
self._load_corpus()
logger.info("Loaded %d %s Documents.", len(self.corpus), split.upper())
logger.info("Doc Example: %s", list(self.corpus.values())[0])
if not len(self.queries):
logger.info("Loading Queries...")
self._load_queries()
if os.path.exists(self.qrels_file):
self._load_qrels()
self.queries = {qid: self.queries[qid] for qid in self.qrels}
logger.info("Loaded %d %s Queries.", len(self.queries), split.upper())
logger.info("Query Example: %s", list(self.queries.values())[0])
return self.corpus, self.queries, self.qrels
def load_corpus(self) -> Dict[str, Dict[str, str]]:
self.check(fIn=self.corpus_file, ext="jsonl")
if not len(self.corpus):
logger.info("Loading Corpus...")
self._load_corpus()
logger.info("Loaded %d Documents.", len(self.corpus))
logger.info("Doc Example: %s", list(self.corpus.values())[0])
return self.corpus
def _load_corpus(self):
num_lines = sum(1 for i in open(self.corpus_file, 'rb'))
with open(self.corpus_file, encoding='utf8') as fIn:
for line in tqdm(fIn, total=num_lines):
line = json.loads(line)
self.corpus[line.get("_id")] = {
"text": line.get("text"),
"title": line.get("title"),
}
def _load_queries(self):
with open(self.query_file, encoding='utf8') as fIn:
for line in fIn:
line = json.loads(line)
self.queries[line.get("_id")] = line.get("text")
def _load_qrels(self):
reader = csv.reader(open(self.qrels_file, encoding="utf-8"),
delimiter="\t", quoting=csv.QUOTE_MINIMAL)
next(reader)
for id, row in enumerate(reader):
query_id, corpus_id, score = row[0], row[1], int(row[2])
if query_id not in self.qrels:
self.qrels[query_id] = {corpus_id: score}
else:
self.qrels[query_id][corpus_id] = score |